Persist the zookeeper client and connection. A step towards a config watcher.

1 files changed, 34 insertions, 23 deletions

diff --git a/config.py b/config.py
index f3a2dfe..b344631 100644
--- a/config.py
+++ b/config.py
@@ -103,6 +103,11 @@ SAVED_MESSAGES: List[str] = []
 PROGRAM_NAME: str = os.path.basename(sys.argv[0])
 ORIG_ARGV: List[str] = sys.argv.copy()
 
+# A zookeeper client that is lazily created so as to not incur the
+# latency of connecting to zookeeper for programs that are not reading
+# or writing their config data into zookeeper.
+ZK: Optional[KazooClient] = None
+
 
 class OptionalRawFormatter(argparse.HelpFormatter):
     """This formatter has the same bahavior as the normal argparse text
@@ -325,6 +330,7 @@ def _augment_sys_argv_from_environment_variables():
 def _augment_sys_argv_from_loadfile():
     """Internal.  Augment with arguments persisted in a saved file."""
 
+    global ZK
     loadfile = None
     saw_other_args = False
     grab_next_arg = False
@@ -344,20 +350,21 @@ def _augment_sys_argv_from_loadfile():
         zkpath = None
         if loadfile[:3] == 'zk:':
             try:
-                zk = KazooClient(
-                    hosts=scott_secrets.ZOOKEEPER_NODES,
-                    use_ssl=True,
-                    verify_certs=False,
-                    keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
-                    keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
-                    certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
-                )
-                zk.start()
+                if ZK is None:
+                    ZK = KazooClient(
+                        hosts=scott_secrets.ZOOKEEPER_NODES,
+                        use_ssl=True,
+                        verify_certs=False,
+                        keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
+                        keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
+                        certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
+                    )
+                    ZK.start()
                 zkpath = loadfile[3:]
                 if not zkpath.startswith('/config/'):
                     zkpath = '/config/' + zkpath
                     zkpath = re.sub(r'//+', '/', zkpath)
-                if not zk.exists(zkpath):
+                if not ZK.exists(zkpath):
                     raise Exception(
                         f'ERROR: --config_loadfile argument must be a file, {loadfile} not found (in zookeeper)'
                     )
@@ -380,7 +387,8 @@ def _augment_sys_argv_from_loadfile():
         newargs = []
         if zkpath:
             try:
-                contents = zk.get(zkpath)[0]
+                assert ZK
+                contents = ZK.get(zkpath)[0]
                 contents = contents.decode()
                 newargs = [
                     arg.strip('\n') for arg in contents.split('\n') if 'config_savefile' not in arg
@@ -404,6 +412,7 @@ def parse(entry_module: Optional[str]) -> Dict[str, Any]:
     global CONFIG_PARSE_CALLED
     if CONFIG_PARSE_CALLED:
         return config
+    global ZK
 
     # If we're about to do the usage message dump, put the main
     # module's argument group last in the list (if possible) so that
@@ -452,18 +461,20 @@ def parse(entry_module: Optional[str]) -> Dict[str, Any]:
                 zkpath = '/config/' + zkpath
                 zkpath = re.sub(r'//+', '/', zkpath)
             try:
-                zk = KazooClient(
-                    hosts=scott_secrets.ZOOKEEPER_NODES,
-                    use_ssl=True,
-                    verify_certs=False,
-                    keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
-                    keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
-                    certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
-                )
-                zk.start()
-                if zk.exists(zkpath):
-                    zk.delete(zkpath)
-                zk.create(zkpath, data.encode())
+                if not ZK:
+                    ZK = KazooClient(
+                        hosts=scott_secrets.ZOOKEEPER_NODES,
+                        use_ssl=True,
+                        verify_certs=False,
+                        keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
+                        keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
+                        certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
+                    )
+                    ZK.start()
+                if not ZK.exists(zkpath):
+                    ZK.create(zkpath, data.encode())
+                else:
+                    ZK.set(zkpath, data.encode())
             except Exception as e:
                 raise Exception(f'Failed to create zookeeper path {zkpath}') from e
             SAVED_MESSAGES.append(f'Saved config to zookeeper in {zkpath}')