From 2ee0b3a0792a0ec904504b0b9d0cae7d9a7cfa3a Mon Sep 17 00:00:00 2001
From: Tim Abbott <tabbott@zulip.com>
Date: Fri, 23 Aug 2013 14:13:26 -0400
Subject: [PATCH] zephyr_mirror: Add options to save/restore sessions.

The davidben-patched-for-roost Zephyr branch (available at
https://github.com/davidben/zephyr/tree/roost) adds Zephyr support for
these options.  We also patch python-zephyr to expose them.  These
basically let you save your Zephyr tickets and port number to a file,
so that you can later restore them (even potentially after the machine
rebooted).  Basically because Zephyr is UDP, the Zephyr server will
continue trying to deliver messages to a particular port number that
was registered for up to 20 minutes after getting an error; so we can
even have downtime and reboot and still get our packets so long as we
restore the sessions within 20 minutes.

(imported from commit 986cbb157ddfa57aa4b644cd826f8418e9876dc7)
---
 bots/zephyr_mirror_backend.py | 51 ++++++++++++++++++++++++++++-------
 1 file changed, 41 insertions(+), 10 deletions(-)

diff --git a/bots/zephyr_mirror_backend.py b/bots/zephyr_mirror_backend.py
index bd04bf9..4dbfefc 100755
--- a/bots/zephyr_mirror_backend.py
+++ b/bots/zephyr_mirror_backend.py
@@ -473,6 +473,21 @@ def zephyr_init_autoretry():
 
     quit_failed_initialization("Could not initialize Zephyr library, quitting!")
 
+def zephyr_load_session_autoretry(session_path):
+    backoff = RandomExponentialBackoff()
+    while backoff.keep_going():
+        try:
+            session = file(session_path, "r").read()
+            zephyr._z.initialize()
+            zephyr._z.load_session(session)
+            zephyr.__inited = True
+            return
+        except IOError:
+            logger.exception("Error loading saved Zephyr session (retrying).  Traceback:")
+            backoff.fail()
+
+    quit_failed_initialization("Could not load saved Zephyr session, quitting!")
+
 def zephyr_subscribe_autoretry(sub):
     backoff = RandomExponentialBackoff()
     while backoff.keep_going():
@@ -489,15 +504,21 @@ def zephyr_subscribe_autoretry(sub):
     quit_failed_initialization("Could not subscribe to personals, quitting!")
 
 def zephyr_to_zulip(options):
-    zephyr_init_autoretry()
-    if options.forward_class_messages:
-        update_subscriptions()
-    if options.forward_personals:
-        # Subscribe to personals; we really can't operate without
-        # those subscriptions, so just retry until it works.
-        zephyr_subscribe_autoretry(("message", "*", "%me%"))
-        if subscribed_to_mail_messages():
-            zephyr_subscribe_autoretry(("mail", "inbox", "%me%"))
+    if options.use_sessions and os.path.exists(options.session_path):
+        logger.info("Loading old session")
+        zephyr_load_session_autoretry(options.session_path)
+    else:
+        zephyr_init_autoretry()
+        if options.forward_class_messages:
+            update_subscriptions()
+        if options.forward_personals:
+            # Subscribe to personals; we really can't operate without
+            # those subscriptions, so just retry until it works.
+            zephyr_subscribe_autoretry(("message", "*", "%me%"))
+            if subscribed_to_mail_messages():
+                zephyr_subscribe_autoretry(("mail", "inbox", "%me%"))
+        if options.use_sessions:
+            file(options.session_path, "w").write(zephyr._z.dump_session())
 
     if options.resend_log_path is not None:
         with open(options.resend_log_path, 'r') as log:
@@ -932,6 +953,13 @@ def parse_args():
     parser.add_option('--root-path',
                       default="/afs/athena.mit.edu/user/t/a/tabbott/for_friends",
                       help=optparse.SUPPRESS_HELP)
+    parser.add_option('--session-path',
+                      default=None,
+                      help=optparse.SUPPRESS_HELP)
+    parser.add_option('--use-sessions',
+                      default=False,
+                      action='store_true',
+                      help=optparse.SUPPRESS_HELP)
     parser.add_option('--test-mode',
                       default=False,
                       help=optparse.SUPPRESS_HELP,
@@ -945,7 +973,7 @@ def die_gracefully(signal, frame):
         # this is a child process, so we want os._exit (no clean-up necessary)
         os._exit(1)
 
-    if CURRENT_STATE == States.ZephyrToZulip:
+    if CURRENT_STATE == States.ZephyrToZulip and not options.use_sessions:
         try:
             # zephyr=>zulip processes may have added subs, so run cancelSubs
             zephyr._z.cancelSubs()
@@ -1040,6 +1068,9 @@ or specify the --api-key-file option.""" % (options.api_key_file,))))
         options.forward_personals = False
         options.forward_from_zulip = False
 
+    if options.session_path is None:
+        options.session_path = "/var/tmp/%s" % (options.user,)
+
     if options.forward_from_zulip:
         child_pid = os.fork()
         if child_pid == 0: