Updated iter_bucket to use concurrent futures.

This commit addresses issue #340.
AWS Lambda environments do not support multiprocessing.Queue or
multiprocessing.Pool, which are used by iter_bucket to optimize the
pulling of files from s3.

Solution: Switch to using concurrent.futures.ThreadPoolExecutor instead.
This still optimizes the pulling of files from s3 without using new
processes.

Author: derpferd

setup.py

@@ -41,7 +41,7 @@ def read(fname):
 
 tests_require = [
     'mock',
-    'moto==1.3.4',
+    'moto[server]==1.3.4',
     'pathlib2',
     'responses',
     # Temporary pin boto3 & botocore, because moto doesn't work with new version

smart_open/s3.py

@@ -21,6 +21,16 @@
 
 logger = logging.getLogger(__name__)
 
+# AWS Lambda environments do not support multiprocessing.Queue or multiprocessing.Pool.
+# However they do support Threads and therefore concurrent.futures's ThreadPoolExecutor.
+# We use this flag to allow python 2 backward compatibility, where concurrent.futures doesn't exist.
+_CONCURRENT_FUTURES = False
+try:
+    import concurrent.futures
+    _CONCURRENT_FUTURES = True
+except ImportError:
+    warnings.warn("concurrent.futures could not be imported and won't be used")
+
 # Multiprocessing is unavailable in App Engine (and possibly other sandboxes).
 # The only method currently relying on it is iter_bucket, which is instructed
 # whether to use it by the MULTIPROCESSING flag.
@@ -766,10 +776,26 @@ def terminate(self):
         pass
 
 
+class ConcurrentFuturesPool(object):
+    """A class that mimics multiprocessing.pool.Pool but uses concurrent futures instead of processes."""
+    def __init__(self, max_workers):
+        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers)
+
+    def imap_unordered(self, function, items):
+        futures = [self.executor.submit(function, item) for item in items]
+        return map(lambda future: future.result(), concurrent.futures.as_completed(futures))
+
+    def terminate(self):
+        self.executor.shutdown(wait=True)
+
+
 @contextlib.contextmanager
 def _create_process_pool(processes=1):
-    if _MULTIPROCESSING and processes:
-        logger.info("creating pool with %i workers", processes)
+    if _CONCURRENT_FUTURES and processes:
+        logger.info("creating concurrent futures pool with %i workers", processes)
+        pool = ConcurrentFuturesPool(max_workers=processes)
+    elif _MULTIPROCESSING and processes:
+        logger.info("creating multiprocessing pool with %i workers", processes)
         pool = multiprocessing.pool.Pool(processes=processes)
     else:
         logger.info("creating dummy pool")

smart_open/tests/test_s3.py

@@ -502,13 +502,89 @@ def test_old(self):
 @maybe_mock_s3
 class IterBucketSingleProcessTest(unittest.TestCase):
     def setUp(self):
-        self.old_flag = smart_open.s3._MULTIPROCESSING
+        self.old_flag_multi = smart_open.s3._MULTIPROCESSING
+        # self.old_flag_concurrent = smart_open.s3._CONCURRENT_FUTURES
         smart_open.s3._MULTIPROCESSING = False
+        # smart_open.s3._CONCURRENT_FUTURES = False
 
         ignore_resource_warnings()
 
     def tearDown(self):
-        smart_open.s3._MULTIPROCESSING = self.old_flag
+        smart_open.s3._MULTIPROCESSING = self.old_flag_multi
+        # smart_open.s3._CONCURRENT_FUTURES = self.old_flag_concurrent
+        cleanup_bucket()
+
+    def test(self):
+        num_keys = 101
+        populate_bucket(num_keys=num_keys)
+        keys = list(smart_open.s3.iter_bucket(BUCKET_NAME))
+        self.assertEqual(len(keys), num_keys)
+
+        expected = [('key_%d' % x, b'%d' % x) for x in range(num_keys)]
+        self.assertEqual(sorted(keys), sorted(expected))
+
+
+@maybe_mock_s3
+class IterBucketConcurrentFuturesTest(unittest.TestCase):
+    def setUp(self):
+        self.old_flag_multi = smart_open.s3._MULTIPROCESSING
+        smart_open.s3._MULTIPROCESSING = False
+        if not smart_open.s3._CONCURRENT_FUTURES:
+            self.fail('Unable to test iter bucket concurrent futures since it is disabled on this machine.')
+
+        ignore_resource_warnings()
+
+    def tearDown(self):
+        smart_open.s3._MULTIPROCESSING = self.old_flag_multi
+        cleanup_bucket()
+
+    def test(self):
+        num_keys = 101
+        populate_bucket(num_keys=num_keys)
+        keys = list(smart_open.s3.iter_bucket(BUCKET_NAME))
+        self.assertEqual(len(keys), num_keys)
+
+        expected = [('key_%d' % x, b'%d' % x) for x in range(num_keys)]
+        self.assertEqual(sorted(keys), sorted(expected))
+
+
+@maybe_mock_s3
+class IterBucketMultiprocessingTest(unittest.TestCase):
+    def setUp(self):
+        self.old_flag_concurrent = smart_open.s3._CONCURRENT_FUTURES
+        smart_open.s3._CONCURRENT_FUTURES = False
+        if not smart_open.s3._MULTIPROCESSING:
+            self.fail('Unable to test iter bucket multiprocessing since it is disabled on this machine.')
+
+        ignore_resource_warnings()
+
+    def tearDown(self):
+        smart_open.s3._CONCURRENT_FUTURES = self.old_flag_concurrent
+        cleanup_bucket()
+
+    def test(self):
+        num_keys = 101
+        populate_bucket(num_keys=num_keys)
+        keys = list(smart_open.s3.iter_bucket(BUCKET_NAME))
+        self.assertEqual(len(keys), num_keys)
+
+        expected = [('key_%d' % x, b'%d' % x) for x in range(num_keys)]
+        self.assertEqual(sorted(keys), sorted(expected))
+
+
+@maybe_mock_s3
+class IterBucketSingleProcessTest(unittest.TestCase):
+    def setUp(self):
+        self.old_flag_multi = smart_open.s3._MULTIPROCESSING
+        self.old_flag_concurrent = smart_open.s3._CONCURRENT_FUTURES
+        smart_open.s3._MULTIPROCESSING = False
+        smart_open.s3._CONCURRENT_FUTURES = False
+
+        ignore_resource_warnings()
+
+    def tearDown(self):
+        smart_open.s3._MULTIPROCESSING = self.old_flag_multi
+        smart_open.s3._CONCURRENT_FUTURES = self.old_flag_concurrent
         cleanup_bucket()
 
     def test(self):