Repository: Backblaze/B2_Command_Line_Tool

Issue:
Using a lot of threads with sync cause sqlite errors
I'm attempting to use the sync command to upload 1 million files using 200 threads. In the log, I'm seeing database errors that prevents an individual file from being uploaded:

`b2_upload(/Volumes/Library/Data/1m-files/file100352.txt, file100352.txt, 1463692063000): OperationalError('database is locked',) database is locked`

When that thread fails - the file is left out of the sync. It never is re-tried.

In addition, infrequently B2 CLI crashes out completely because auth credentials couldn't be retrieved:

`ERROR: Missing account data: database is locked  Use: b2 authorize_account`



Comments:

Patch:
diff --git a/b2/account_info.py b/b2/account_info.py
index 9a639c5..5eafba6 100644
--- a/b2/account_info.py
+++ b/b2/account_info.py
@@ -141,6 +141,42 @@ class AbstractAccountInfo(object):
         pass
 
 
+class UploadUrlPool(object):
+    """
+    For each key (either a bucket id or large file id), holds a pool
+    of (url, auth_token) pairs, with thread-safe methods to add and
+    remove them.
+    """
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._pool = collections.defaultdict(list)
+
+    def put(self, key, url, auth_token):
+        """
+        Adds the url and auth token to the pool for the given key.
+        """
+        with self._lock:
+            pair = (url, auth_token)
+            self._pool[key].append(pair)
+
+    def take(self, key):
+        """
+        Returns (url, auth_token) if one is available, or (None, None) if not.
+        """
+        with self._lock:
+            pair_list = self._pool[key]
+            if pair_list:
+                return pair_list.pop()
+            else:
+                return (None, None)
+
+    def clear_for_key(self, key):
+        with self._lock:
+            if key in self._pool:
+                del self._pool[key]
+
+
 class SqliteAccountInfo(AbstractAccountInfo):
     """
     Stores account information in an sqlite database, which is
@@ -157,9 +193,8 @@ class SqliteAccountInfo(AbstractAccountInfo):
         with self._get_connection() as conn:
             self._create_tables(conn)
 
-        self._large_file_uploads = collections.defaultdict(
-            list
-        )  # We don't keep large file upload URLs across a reload
+        self._bucket_uploads = UploadUrlPool()
+        self._large_file_uploads = UploadUrlPool()
 
         # this lock controls access to self._large_file_uploads
         self._lock = threading.Lock()
@@ -267,6 +302,9 @@ class SqliteAccountInfo(AbstractAccountInfo):
            );
         """
         )
+        # This table is not used any more.  We may use it again
+        # someday if we save upload URLs across invocations of
+        # the command-line tool.
         conn.execute(
             """
            CREATE TABLE IF NOT EXISTS
@@ -367,48 +405,22 @@ class SqliteAccountInfo(AbstractAccountInfo):
             return None
 
     def put_bucket_upload_url(self, bucket_id, upload_url, upload_auth_token):
-        with self._get_connection() as conn:
-            conn.execute(
-                'INSERT INTO bucket_upload_url (bucket_id, upload_url, upload_auth_token) values (?, ?, ?);',
-                (bucket_id, upload_url, upload_auth_token)
-            )
+        self._bucket_uploads.put(bucket_id, upload_url, upload_auth_token)
 
     def clear_bucket_upload_data(self, bucket_id):
-        with self._get_connection() as conn:
-            conn.execute('DELETE FROM bucket_upload_url WHERE bucket_id = ?;', (bucket_id,))
+        self._bucket_uploads.clear_for_key(bucket_id)
 
     def take_bucket_upload_url(self, bucket_id):
-        try:
-            with self._get_connection() as conn:
-                cursor = conn.execute(
-                    'SELECT upload_url, upload_auth_token FROM bucket_upload_url WHERE bucket_id = ?;',
-                    (bucket_id,)
-                )
-                (upload_url, upload_auth_token) = cursor.fetchone()
-                conn.execute(
-                    'DELETE FROM bucket_upload_url WHERE upload_auth_token = ?;',
-                    (upload_auth_token,)
-                )
-                return (upload_url, upload_auth_token)
-        except:
-            return (None, None)
+        return self._bucket_uploads.take(bucket_id)
 
     def put_large_file_upload_url(self, file_id, upload_url, upload_auth_token):
-        with self._lock:
-            self._large_file_uploads[file_id].append((upload_url, upload_auth_token))
+        self._large_file_uploads.put(file_id, upload_url, upload_auth_token)
 
     def take_large_file_upload_url(self, file_id):
-        with self._lock:
-            url_list = self._large_file_uploads.get(file_id, [])
-            if len(url_list) == 0:
-                return (None, None)
-            else:
-                return url_list.pop()
+        return self._large_file_uploads.take(file_id)
 
     def clear_large_file_upload_urls(self, file_id):
-        with self._lock:
-            if file_id in self._large_file_uploads:
-                del self._large_file_uploads[file_id]
+        self._large_file_uploads.clear_for_key(file_id)
 
 
 class StubAccountInfo(AbstractAccountInfo):


Test patch:
diff --git a/test/test_account_info.py b/test/test_account_info.py
index 2a52183..68a6f22 100644
--- a/test/test_account_info.py
+++ b/test/test_account_info.py
@@ -14,7 +14,7 @@ import unittest
 
 import six
 
-from b2.account_info import SqliteAccountInfo
+from b2.account_info import SqliteAccountInfo, UploadUrlPool
 from b2.exception import CorruptAccountInfo, MissingAccountData
 
 try:
@@ -23,6 +23,32 @@ except:
     import mock
 
 
+class TestUploadUrlPool(unittest.TestCase):
+    def setUp(self):
+        self.pool = UploadUrlPool()
+
+    def test_take_empty(self):
+        self.assertEqual((None, None), self.pool.take('a'))
+
+    def test_put_and_take(self):
+        self.pool.put('a', 'url_a1', 'auth_token_a1')
+        self.pool.put('a', 'url_a2', 'auth_token_a2')
+        self.pool.put('b', 'url_b1', 'auth_token_b1')
+        self.assertEqual(('url_a2', 'auth_token_a2'), self.pool.take('a'))
+        self.assertEqual(('url_a1', 'auth_token_a1'), self.pool.take('a'))
+        self.assertEqual((None, None), self.pool.take('a'))
+        self.assertEqual(('url_b1', 'auth_token_b1'), self.pool.take('b'))
+        self.assertEqual((None, None), self.pool.take('b'))
+
+    def test_clear(self):
+        self.pool.put('a', 'url_a1', 'auth_token_a1')
+        self.pool.clear_for_key('a')
+        self.pool.put('b', 'url_b1', 'auth_token_b1')
+        self.assertEqual((None, None), self.pool.take('a'))
+        self.assertEqual(('url_b1', 'auth_token_b1'), self.pool.take('b'))
+        self.assertEqual((None, None), self.pool.take('b'))
+
+
 class TestSqliteAccountInfo(unittest.TestCase):
 
     FILE_NAME = '/tmp/test_b2_account_info'
@@ -99,20 +125,6 @@ class TestSqliteAccountInfo(unittest.TestCase):
         except MissingAccountData:
             pass
 
-    def test_bucket_upload_data(self):
-        account_info = self._make_info()
-        account_info.put_bucket_upload_url('bucket-0', 'http://bucket-0', 'bucket-0_auth')
-        self.assertEqual(
-            ('http://bucket-0', 'bucket-0_auth'), account_info.take_bucket_upload_url('bucket-0')
-        )
-        self.assertEqual((None, None), self._make_info().take_bucket_upload_url('bucket-0'))
-        account_info.put_bucket_upload_url('bucket-0', 'http://bucket-0', 'bucket-0_auth')
-        self.assertEqual(
-            ('http://bucket-0', 'bucket-0_auth'),
-            self._make_info().take_bucket_upload_url('bucket-0')
-        )
-        self.assertEqual((None, None), account_info.take_bucket_upload_url('bucket-0'))
-
     def test_clear_bucket_upload_data(self):
         account_info = self._make_info()
         account_info.put_bucket_upload_url('bucket-0', 'http://bucket-0', 'bucket-0_auth')