Merge pull request #352 from Backblaze/max-streams

Expose max streams parameter

Author: ppolewicz

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 * Logging performance summary of parallel download threads
+* Add `max_download_streams_per_file` parameter to B2Api class and underlying structures
 
 ### Fixed
 * Replace `ReplicationScanResult.source_has_sse_c_enabled` with `source_encryption_mode`

b2sdk/api.py

@@ -64,6 +64,7 @@ def __init__(
         max_download_workers: Optional[int] = None,
         save_to_buffer_size: Optional[int] = None,
         check_download_hash: bool = True,
+        max_download_streams_per_file: Optional[int] = None,
     ):
         """
         Initialize Services object using given session.
@@ -74,6 +75,7 @@ def __init__(
         :param max_download_workers: maximum number of download threads
         :param save_to_buffer_size: buffer size to use when writing files using DownloadedFile.save_to
         :param check_download_hash: whether to check hash of downloaded files. Can be disabled for files with internal checksums, for example, or to forcefully retrieve objects with corrupted payload or hash value
+        :param max_download_streams_per_file: how many streams to use for parallel downloader
         """
         self.api = api
         self.session = api.session
@@ -82,11 +84,13 @@ def __init__(
             services=self, max_workers=max_upload_workers
         )
         self.copy_manager = self.COPY_MANAGER_CLASS(services=self, max_workers=max_copy_workers)
+        assert max_download_streams_per_file is None or max_download_streams_per_file >= 1
         self.download_manager = self.DOWNLOAD_MANAGER_CLASS(
             services=self,
             max_workers=max_download_workers,
             write_buffer_size=save_to_buffer_size,
             check_hash=check_download_hash,
+            max_download_streams_per_file=max_download_streams_per_file,
         )
         self.emerger = Emerger(self)
 
@@ -129,6 +133,7 @@ def __init__(
         max_download_workers: Optional[int] = None,
         save_to_buffer_size: Optional[int] = None,
         check_download_hash: bool = True,
+        max_download_streams_per_file: Optional[int] = None,
     ):
         """
         Initialize the API using the given account info.
@@ -145,6 +150,7 @@ def __init__(
         :param max_download_workers: maximum number of download threads
         :param save_to_buffer_size: buffer size to use when writing files using DownloadedFile.save_to
         :param check_download_hash: whether to check hash of downloaded files. Can be disabled for files with internal checksums, for example, or to forcefully retrieve objects with corrupted payload or hash value
+        :param max_download_streams_per_file: number of streams for parallel download manager
         """
         self.session = self.SESSION_CLASS(
             account_info=account_info, cache=cache, api_config=api_config
@@ -159,6 +165,7 @@ def __init__(
             max_download_workers=max_download_workers,
             save_to_buffer_size=save_to_buffer_size,
             check_download_hash=check_download_hash,
+            max_download_streams_per_file=max_download_streams_per_file,
         )
 
     @property

b2sdk/transfer/inbound/download_manager.py

@@ -44,7 +44,13 @@ class DownloadManager(TransferManager, ThreadPoolMixin, metaclass=B2TraceMetaAbs
     PARALLEL_DOWNLOADER_CLASS = staticmethod(ParallelDownloader)
     SIMPLE_DOWNLOADER_CLASS = staticmethod(SimpleDownloader)
 
-    def __init__(self, write_buffer_size: Optional[int] = None, check_hash: bool = True, **kwargs):
+    def __init__(
+        self,
+        write_buffer_size: Optional[int] = None,
+        check_hash: bool = True,
+        max_download_streams_per_file: Optional[int] = None,
+        **kwargs
+    ):
         """
         Initialize the DownloadManager using the given services object.
         """
@@ -58,6 +64,7 @@ def __init__(self, write_buffer_size: Optional[int] = None, check_hash: bool = T
                 align_factor=write_buffer_size,
                 thread_pool=self._thread_pool,
                 check_hash=check_hash,
+                max_streams=max_download_streams_per_file,
             ),
             self.SIMPLE_DOWNLOADER_CLASS(
                 min_chunk_size=self.MIN_CHUNK_SIZE,

test/unit/v_all/test_api.py

@@ -22,6 +22,86 @@
 from ..test_base import TestBase
 
 
+class DummyA:
+    def __init__(self, *args, **kwargs):
+        pass
+
+
+class DummyB:
+    def __init__(self, *args, **kwargs):
+        pass
+
+
+class TestServices:
+    @pytest.mark.apiver(from_ver=2)
+    @pytest.mark.parametrize(
+        ('kwargs', '_raw_api_class'),
+        [
+            [
+                {
+                    'max_upload_workers': 1,
+                    'max_copy_workers': 2,
+                    'max_download_workers': 3,
+                    'save_to_buffer_size': 4,
+                    'check_download_hash': False,
+                    'max_download_streams_per_file': 5,
+                },
+                DummyA,
+            ],
+            [
+                {
+                    'max_upload_workers': 2,
+                    'max_copy_workers': 3,
+                    'max_download_workers': 4,
+                    'save_to_buffer_size': 5,
+                    'check_download_hash': True,
+                    'max_download_streams_per_file': 6,
+                },
+                DummyB,
+            ],
+        ],
+    )  # yapf: disable
+    def test_api_initialization(self, kwargs, _raw_api_class):
+        self.account_info = InMemoryAccountInfo()
+        self.cache = InMemoryCache()
+
+        api_config = B2HttpApiConfig(_raw_api_class=_raw_api_class)
+
+        self.api = B2Api(
+            self.account_info,
+            self.cache,
+            api_config=api_config,
+
+            **kwargs
+        )  # yapf: disable
+
+        assert self.api.account_info is self.account_info
+        assert self.api.api_config is api_config
+        assert self.api.cache is self.cache
+
+        assert self.api.session.account_info is self.account_info
+        assert self.api.session.cache is self.cache
+        assert isinstance(self.api.session.raw_api, _raw_api_class)
+
+        assert isinstance(self.api.file_version_factory, B2Api.FILE_VERSION_FACTORY_CLASS)
+        assert isinstance(
+            self.api.download_version_factory,
+            B2Api.DOWNLOAD_VERSION_FACTORY_CLASS,
+        )
+
+        services = self.api.services
+        assert isinstance(services, B2Api.SERVICES_CLASS)
+
+        # max copy/upload/download workers could only be verified with mocking
+
+        download_manager = services.download_manager
+        assert isinstance(download_manager, services.DOWNLOAD_MANAGER_CLASS)
+
+        assert download_manager.write_buffer_size == kwargs['save_to_buffer_size']
+        assert download_manager.check_hash == kwargs['check_download_hash']
+        assert download_manager.strategies[0].max_streams == kwargs['max_download_streams_per_file']
+
+
 class TestApi(TestBase):
     def setUp(self):
         self.account_info = InMemoryAccountInfo()