Use multiprocessing.Pool callback= to track thread state

Author: timvaillancourt

mongodb_consistent_backup/Archive/Tar/Tar.py

@@ -3,6 +3,7 @@
 
 from copy_reg import pickle
 from multiprocessing import Pool, TimeoutError
+from time import sleep
 from types import MethodType
 
 from TarThread import TarThread
@@ -30,20 +31,23 @@ def __init__(self, manager, config, timer, base_dir, backup_dir, **kwargs):
         self._pool   = None
         self._pooled = []
 
+    def done(self, done_dir):
+        if done_dir in self._pooled:
+            logging.debug("Archiving completed for: %s" % done_dir)
+            self._pooled.remove(done_dir)
+        else:
+            raise OperationError("Unexpected response from tar thread: %s" % done_dir)
+
     def wait(self):
         if len(self._pooled) > 0:
             self._pool.close()
-            logging.debug("Waiting for tar threads to stop")
-            while len(self._pooled) > 0:
-                try:
-                    item = self._pooled[0]
-                    path, result = item
-                    result.get(1)
-                    logging.debug("Archiving completed for directory: %s" % path)
-                    self._pooled.remove(item)
-                except TimeoutError:
-                    continue
+            while len(self._pooled):
+                logging.debug("Waiting for %i tar thread(s) to stop" % len(self._pooled))
+                sleep(2)
+            self._pool.terminate()
+            logging.debug("Stopped all tar threads")
             self.stopped = True
+            self.running = False
 
     def run(self):
         try:
@@ -64,20 +68,18 @@ def run(self):
                     output_file = "%s.tar" % subdir_name
                     if self.do_gzip():
                         output_file  = "%s.tgz" % subdir_name
-                    result = self._pool.apply_async(TarThread(subdir_name, output_file, self.do_gzip(), self.verbose, self.binary).run)
-                    self._pooled.append((subdir_name, result))
+                    self._pool.apply_async(TarThread(subdir_name, output_file, self.do_gzip(), self.verbose, self.binary).run, callback=self.done)
+                    self._pooled.append(subdir_name)
             except Exception, e:
                 self._pool.terminate()
                 logging.fatal("Could not create tar archiving thread! Error: %s" % e)
                 raise Error(e)
-            finally:
-                self.wait()
+            self.wait()
             self.completed = True
 
     def close(self, code=None, frame=None):
-        logging.debug("Stopping tar archiving threads")
         if not self.stopped and self._pool is not None:
+            logging.debug("Stopping tar archiving threads")
             self._pool.terminate()
-            self._pool.join()
             logging.info("Stopped all tar archiving threads")
             self.stopped = True

mongodb_consistent_backup/Archive/Tar/TarThread.py

@@ -50,3 +50,4 @@ def run(self):
             else:
                 logging.fatal("Output file: %s already exists!" % self.output_file)
                 sys.exit(1)
+            return self.backup_dir

mongodb_consistent_backup/Oplog/Resolver/Resolver.py

@@ -4,7 +4,8 @@
 # noinspection PyPackageRequirements
 from bson.timestamp import Timestamp
 from copy_reg import pickle
-from multiprocessing import Pool
+from multiprocessing import Pool, TimeoutError
+from time import sleep
 from types import MethodType
 
 from ResolverThread import ResolverThread
@@ -34,16 +35,20 @@ def __init__(self, manager, config, timer, base_dir, backup_dir, tailed_oplogs,
         self.resolver_summary      = {}
         self.resolver_state        = {}
 
+        self._pool   = None
+        self._pooled = []
         try:
             self._pool = Pool(processes=self.threads(None, 2))
         except Exception, e:
             logging.fatal("Could not start oplog resolver pool! Error: %s" % e)
             raise Error(e)
 
-    def close(self):
-       if self._pool:
+    def close(self, code=None, frame=None):
+       if self._pool and not self.stopped:
+           logging.debug("Stopping all oplog resolver threads")
            self._pool.terminate()
-           self._pool.join()
+           logging.info("Stopped all oplog resolver threads")
+           self.stopped = True
 
     def get_consistent_end_ts(self):
         ts = None
@@ -54,8 +59,26 @@ def get_consistent_end_ts(self):
                     ts = Timestamp(instance['last_ts'].time, 0)
         return ts
 
+    def done(self, done_uri):
+        if done_uri in self._pooled:
+            logging.debug("Resolving completed for: %s" % done_uri)
+            self._pooled.remove(done_uri)
+        else:
+            raise OperationError("Unexpected response from resolver thread: %s" % done_uri)
+
+    def wait(self):
+        if len(self._pooled) > 0:
+            self._pool.close()
+            while len(self._pooled):
+                logging.debug("Waiting for %i oplog resolver thread(s) to stop" % len(self._pooled))
+                sleep(2)
+            self._pool.terminate()
+            logging.debug("Stopped all oplog resolve threads")
+            self.stopped = True
+            self.running = False
+
     def run(self):
-        logging.info("Resolving oplogs (options: threads=%s,compression=%s)" % (self.threads(), self.compression()))
+        logging.info("Resolving oplogs (options: threads=%s, compression=%s)" % (self.threads(), self.compression()))
         self.timer.start(self.timer_name)
 
         for shard in self.backup_oplogs:
@@ -80,14 +103,15 @@ def run(self):
                             backup_oplog.copy(),
                             self.get_consistent_end_ts(),
                             self.do_gzip()
-                        ).run)
+                        ).run, callback=self.done)
+                        self._pooled.append(uri.str())
                     except Exception, e:
                         logging.fatal("Resolve failed for %s! Error: %s" % (uri, e))
                         raise Error(e)
             else:
                 logging.info("No tailed oplog for host %s" % uri)
-        self._pool.close()
-        self._pool.join()
+        self.wait()
+        self.completed = True
 
         self.timer.stop(self.timer_name)
         logging.info("Oplog resolving completed in %.2f seconds" % self.timer.duration(self.timer_name))

mongodb_consistent_backup/Oplog/Resolver/ResolverThread.py

@@ -18,13 +18,7 @@ def __init__(self, state, uri, tailed_oplog, mongodump_oplog, max_end_ts, dump_g
 
         self.oplogs  = {}
         self.changes = 0
-
-    def cleanup(self):
-        if 'tailed' in self.oplogs:
-            self.oplogs['tailed'].close()
-            del self.oplogs['tailed']
-        if 'file' in self.tailed_oplog and os.path.isfile(self.tailed_oplog['file']):
-            os.remove(self.tailed_oplog['file'])
+        self.stopped = True
 
     def run(self):
         self.oplogs['backup'] = Oplog(self.mongodump_oplog['file'], self.dump_gzip, 'a+')
@@ -55,12 +49,16 @@ def run(self):
             self.close()
 
         if self.exit_code == 0:
-            logging.info("Applied %i oplog changes to %s oplog, end ts: %s" % (self.changes, self.uri, self.mongodump_oplog_h.last_ts()))
+            logging.info("Applied %i oplog changes to %s oplog, end ts: %s" % (self.changes, self.uri, self.last_ts))
+            return self.uri.str()
 
     def close(self):
-        self.cleanup()
-        if len(self.oplogs) > 0:
+        if len(self.oplogs) > 0 and not self.stopped:
+            logging.debug("Closing oplog file handles")
             for oplog in self.oplogs:
                 self.oplogs[oplog].flush()
                 self.oplogs[oplog].close()
-                del self.oplogs[oplog]
+            self.stopped = True
+        if 'file' in self.tailed_oplog and os.path.isfile(self.tailed_oplog['file']):
+            logging.debug("Removing temporary/tailed oplog file: %s" % self.tailed_oplog['file'])
+            os.remove(self.tailed_oplog['file'])