[patch] Use concurrent.futures.ProcessPoolExecutor for python3.5+

executor_python3_5.patch (text/plain), 1.94 KB, created by Zac Medico on 2018-03-13 23:52:57 UTC

diff --git a/gemato/recursiveloader.py b/gemato/recursiveloader.py
index 0969554..54dbc9d 100644
--- a/gemato/recursiveloader.py
+++ b/gemato/recursiveloader.py
@@ -8,6 +8,11 @@ import multiprocessing
 import os.path
 import sys
 
+try:
+	from concurrent.futures import ProcessPoolExecutor
+except ImportError:
+	ProcessPoolExecutor = None
+
 import gemato.compression
 import gemato.exceptions
 import gemato.manifest
@@ -16,6 +21,38 @@ import gemato.util
 import gemato.verify
 
 
+class Pool(object):
+	def __init__(self, processes=None):
+		# map supports chunksize in python3.5+
+		self._executor = sys.version_info >= (3, 5)
+		if self._executor:
+			self._pool = ProcessPoolExecutor(max_workers=processes)
+			self.imap_unordered = self._pool.map
+		else:
+			self._pool = multiprocessing.Pool(processes=processes)
+			self.imap_unordered = self._pool.imap_unordered
+
+		self.map = self._pool.map
+
+	def close(self):
+		if self._executor:
+			self._pool.shutdown()
+		else:
+			self._pool.close()
+
+	def join(self):
+		if self._executor:
+			self._pool.shutdown(wait=True)
+		else:
+			self._pool.join()
+
+	def terminate(self):
+		if self._executor:
+			self._pool.shutdown(wait=True)
+		else:
+			self._pool.terminate()
+
+
 class ManifestLoader(object):
     """
     Helper class to load Manifests in subprocesses.
@@ -382,7 +419,7 @@ class ManifestRecursiveLoader(object):
         unconditionally of whether they match parent checksums.
         """
 
-        pool = multiprocessing.Pool(processes=self.max_jobs)
+        pool = Pool(processes=self.max_jobs)
 
         try:
             # TODO: figure out how to avoid confusing uses of 'recursive'
@@ -657,7 +694,7 @@ class ManifestRecursiveLoader(object):
                 self.manifest_device,
                 fail_handler, last_mtime)
 
-        pool = multiprocessing.Pool(processes=self.max_jobs)
+        pool = Pool(processes=self.max_jobs)
 
         try:
             # verify the directories in parallel