Repository: Backblaze/B2_Command_Line_Tool

Issue:
Limit sync action queue length to avoid memory explosion
found in #279 

Comments:

Patch:
diff --git a/README.md b/README.md
index 9ffca80..0548c49 100644
--- a/README.md
+++ b/README.md
@@ -47,6 +47,7 @@ this:
     b2 list-unfinished-large-files <bucketName>
     b2 ls [--long] [--versions] <bucketName> [<folderName>]
     b2 make-url <fileId>
+    b2 show-account-info
     b2 sync [--delete] [--keepDays N] [--skipNewer] [--replaceNewer] \
         [--compareVersions <option>] [--threads N] [--noProgress] \
         [--excludeRegex <regex> [--includeRegex <regex>]] [--dryRun] \
diff --git a/b2/console_tool.py b/b2/console_tool.py
index a657532..b218880 100644
--- a/b2/console_tool.py
+++ b/b2/console_tool.py
@@ -629,6 +629,23 @@ class MakeUrl(Command):
         return 0
 
 
+class ShowAccountInfo(Command):
+    """
+    b2 show-account-info
+
+        Shows the account ID, key, auth token, and URLs.
+    """
+
+    def run(self, args):
+        account_info = self.api.account_info
+        self._print('Account ID:         %s' % (account_info.get_account_id(),))
+        self._print('Application Key:    %s' % (account_info.get_application_key(),))
+        self._print('Account Auth Token: %s' % (account_info.get_account_auth_token(),))
+        self._print('API URL:            %s' % (account_info.get_api_url(),))
+        self._print('Download URL:       %s' % (account_info.get_download_url(),))
+        return 0
+
+
 class Sync(Command):
     """
     b2 sync [--delete] [--keepDays N] [--skipNewer] [--replaceNewer] \\
diff --git a/b2/sync/sync.py b/b2/sync/sync.py
index eaba7a6..9dcf371 100644
--- a/b2/sync/sync.py
+++ b/b2/sync/sync.py
@@ -13,6 +13,7 @@ from __future__ import division
 import logging
 import re
 import six
+import threading
 
 from ..exception import CommandError
 from ..utils import trace_call
@@ -162,6 +163,40 @@ def count_files(local_folder, reporter):
     reporter.end_local()
 
 
+class BoundedQueueExecutor(object):
+    """
+    Wraps a futures.Executor and limits the number of requests that
+    can be queued at once.  Requests to submit() tasks block until
+    there is room in the queue.
+
+    The number of available slots in the queue is tracked with a
+    semaphore that is acquired before queueing an action, and
+    released when an action finishes.
+    """
+
+    def __init__(self, executor, queue_limit):
+        self.executor = executor
+        self.semaphore = threading.Semaphore(queue_limit)
+
+    def submit(self, fcn, *args, **kwargs):
+        # Wait until there is room in the queue.
+        self.semaphore.acquire()
+
+        # Wrap the action in a function that will release
+        # the semaphore after it runs.
+        def run_it():
+            try:
+                fcn(*args, **kwargs)
+            finally:
+                self.semaphore.release()
+
+        # Submit the wrapped action.
+        return self.executor.submit(run_it)
+
+    def shutdown(self):
+        self.executor.shutdown()
+
+
 @trace_call(logger)
 def sync_folders(
     source_folder, dest_folder, args, now_millis, stdout, no_progress, max_workers, dry_run=False
@@ -183,7 +218,12 @@ def sync_folders(
         # not the same as the executor in the API object, which is used for
         # uploads.  The tasks in this executor wait for uploads.  Putting them
         # in the same thread pool could lead to deadlock.
-        sync_executor = futures.ThreadPoolExecutor(max_workers=max_workers)
+        #
+        # We use an executor with a bounded queue to avoid using up lots of memory
+        # when syncing lots of files.
+        unbounded_executor = futures.ThreadPoolExecutor(max_workers=max_workers)
+        queue_limit = max_workers + 1000
+        sync_executor = BoundedQueueExecutor(unbounded_executor, queue_limit=queue_limit)
 
         # First, start the thread that counts the local files.  That's the operation
         # that should be fastest, and it provides scale for the progress reporting.


Test patch:
diff --git a/test/test_console_tool.py b/test/test_console_tool.py
index 99adff5..f046017 100644
--- a/test/test_console_tool.py
+++ b/test/test_console_tool.py
@@ -452,6 +452,17 @@ class TestConsoleTool(TestBase):
                 ], expected_stdout, '', 0
             )
 
+    def test_show_account_info(self):
+        self._authorize_account()
+        expected_stdout = '''
+        Account ID:         my-account
+        Application Key:    good-app-key
+        Account Auth Token: AUTH:my-account
+        API URL:            http://api.example.com
+        Download URL:       http://download.example.com
+        '''
+        self._run_command(['show-account-info'], expected_stdout, '', 0)
+
     def test_sync(self):
         self._authorize_account()
         self._create_my_bucket()
diff --git a/test/test_sync.py b/test/test_sync.py
index 1bf26aa..63dd640 100644
--- a/test/test_sync.py
+++ b/test/test_sync.py
@@ -12,6 +12,8 @@ from __future__ import print_function
 
 import os
 import platform
+import threading
+import time
 import unittest
 
 import six
@@ -21,7 +23,7 @@ from b2.exception import CommandError, DestFileNewer
 from b2.file_version import FileVersionInfo
 from b2.sync.folder import AbstractFolder, B2Folder, LocalFolder
 from b2.sync.file import File, FileVersion
-from b2.sync.sync import make_folder_sync_actions, zip_folders
+from b2.sync.sync import BoundedQueueExecutor, make_folder_sync_actions, zip_folders
 from b2.sync.folder_parser import parse_sync_folder
 from b2.utils import TempDir
 
@@ -30,6 +32,11 @@ try:
 except ImportError:
     from mock import MagicMock
 
+try:
+    import concurrent.futures as futures
+except ImportError:
+    import futures
+
 DAY = 86400000  # milliseconds
 TODAY = DAY * 100  # an arbitrary reference time for testing
 
@@ -689,5 +696,78 @@ class TestMakeSyncActions(TestSync):
         self.assertEqual(expected_actions, [str(a) for a in actions])
 
 
+class TestBoundedQueueExecutor(TestBase):
+    def test_run_more_than_queue_size(self):
+        """
+        Makes sure that the executor will run more jobs that the
+        queue size, which ensures that the semaphore gets released,
+        even if an exception is thrown.
+        """
+        raw_executor = futures.ThreadPoolExecutor(1)
+        bounded_executor = BoundedQueueExecutor(raw_executor, 5)
+
+        class Counter(object):
+            """
+            Counts how many times run() is called.
+            """
+
+            def __init__(self):
+                self.counter = 0
+
+            def run(self):
+                """
+                Always increments the counter.  Sometimes raises an exception.
+                """
+                self.counter += 1
+                if self.counter % 2 == 0:
+                    raise Exception('test')
+
+        counter = Counter()
+        for _ in six.moves.range(10):
+            bounded_executor.submit(counter.run)
+        bounded_executor.shutdown()
+        self.assertEqual(10, counter.counter)
+
+    def test_wait_for_running_jobs(self):
+        """
+        Makes sure that no more than queue_limit workers are
+        running at once, which checks that the semaphore is
+        acquired before submitting an action.
+        """
+        raw_executor = futures.ThreadPoolExecutor(2)
+        bounded_executor = BoundedQueueExecutor(raw_executor, 1)
+        assert_equal = self.assertEqual
+
+        class CountAtOnce(object):
+            """
+            Counts how many threads are running at once.
+            There should never be more than 1 because that's
+            the limit on the bounded executor.
+            """
+
+            def __init__(self):
+                self.running_at_once = 0
+                self.lock = threading.Lock()
+
+            def run(self):
+                with self.lock:
+                    self.running_at_once += 1
+                    assert_equal(1, self.running_at_once)
+                # While we are sleeping here, no other actions should start
+                # running.  If they do, they will increment the counter and
+                # fail the above assertion.
+                time.sleep(0.05)
+                with self.lock:
+                    self.running_at_once -= 1
+                self.counter += 1
+                if self.counter % 2 == 0:
+                    raise Exception('test')
+
+        count_at_once = CountAtOnce()
+        for _ in six.moves.range(5):
+            bounded_executor.submit(count_at_once.run)
+        bounded_executor.shutdown()
+
+
 if __name__ == '__main__':
     unittest.main()