broke replset into more defs for upcoming safe-oplogtail-stop fix, keep pool of DB+Replset objects because they'll be called later

Author: Tim Vaillancourt

MongoBackup/Backup.py

@@ -143,9 +143,6 @@ def cleanup_and_exit(self, code, frame):
         if current_process().name == "MainProcess":
             logging.info("Starting cleanup and exit procedure! Killing running threads")
 
-            if self.sharding:
-                self.sharding.restore_balancer_state()
-
             submodules = ['replset', 'sharding', 'mongodumper', 'oplogtailer', 'archiver', 'uploader_s3']
             for submodule_name in submodules:
                 submodule = getattr(self, submodule_name)

MongoBackup/Replset.py

@@ -143,7 +143,8 @@ def __init__(self, sharding, db, user, password, authdb, max_lag_secs):
         self.authdb       = authdb
         self.max_lag_secs = max_lag_secs
 
-        self.replset = None
+        self.replsets      = {} 
+        self.replset_conns = {}
 
         # Check Sharding class:
         if not self.sharding.__class__.__name__ == "Sharding":
@@ -161,22 +162,41 @@ def __init__(self, sharding, db, user, password, authdb, max_lag_secs):
             logging.fatal("Could not get DB connection! Error: %s" % e)
             raise e
 
+    def get_replset_connection(self, host, port, force=False):
+        conn_name = "%s-%i" % (host, port)
+        if force or not conn_name in self.replset_conns:
+            try:
+                self.replset_conns[conn_name] = DB(host, port, self.user, self.password, self.authdb)
+            except Exception, e:
+                logging.fatal("Could not get DB connection to %s:%i! Error: %s" % (host, port, e))
+                raise e
+        return self.replset_conns[conn_name]
+
+    def get_replsets(self, force=False):
+        for shard in self.sharding.shards():
+            shard_name, members = shard['host'].split('/')
+            host, port = members.split(',')[0].split(":")
+            port       = int(port)
+            if force or not shard_name in self.replsets:
+                try:
+                    rs_db = self.get_replset_connection(host, port)
+                    self.replsets[shard_name] = Replset(rs_db, self.user, self.password, self.authdb, self.max_lag_secs)
+                except Exception, e:
+                    logging.fatal("Could not get Replset class object for replset %s! Error: %s" % (rs_name, e))
+                    raise e
+        print self.replsets
+        return self.replsets
+
     def find_secondaries(self):
         shard_secondaries = {}
-        if self.sharding:
-            for shard in self.sharding.shards():
-                shard_name, members = shard['host'].split('/')
-                host, port          = members.split(',')[0].split(":")
-
-                replset_db   = DB(host, port, self.user, self.password, self.authdb) 
-                self.replset = Replset(replset_db, self.user, self.password, self.authdb, self.max_lag_secs)
-                secondary    = self.replset.find_secondary()
-                shard_secondaries[shard_name] = secondary
-
-                self.replset.close()
-                replset_db.close()
+        for rs_name in self.get_replsets():
+            replset   = self.replsets[rs_name]
+            secondary = replset.find_secondary()
+            shard_secondaries[rs_name] = secondary
         return shard_secondaries
 
     def close(self):
-        if self.replset:
-            self.replset.close()
+        for conn_name in self.replset_conns:
+            self.replset_conns[conn_name].close()
+        for rs_name in self.replsets:
+            self.replsets[rs_name].close()