From ed48c424290428d3acbb6965c3ee1709056b040b Mon Sep 17 00:00:00 2001
From: Maciej Urbanski <maciej.rooter.urbanski@reef.pl>
Date: Mon, 26 Feb 2024 11:37:10 +0100
Subject: [PATCH] fix typing on thread_pool and add tests

---
 b2sdk/utils/thread_pool.py          | 28 +++++++++++++++++--
 test/unit/utils/test_thread_pool.py | 43 +++++++++++++++++++++++++++++
 test/unit/v_all/test_transfer.py    | 18 ++++++------
 3 files changed, 78 insertions(+), 11 deletions(-)
 create mode 100644 test/unit/utils/test_thread_pool.py

diff --git a/b2sdk/utils/thread_pool.py b/b2sdk/utils/thread_pool.py
index 7726cdcb5..b9a190c89 100644
--- a/b2sdk/utils/thread_pool.py
+++ b/b2sdk/utils/thread_pool.py
@@ -9,12 +9,29 @@
 ######################################################################
 from __future__ import annotations
 
+import os
 from concurrent.futures import Future, ThreadPoolExecutor
 from typing import Callable
 
+try:
+    from typing_extensions import Protocol
+except ImportError:
+    from typing import Protocol
+
 from b2sdk.utils import B2TraceMetaAbstract
 
 
+class DynamicThreadPoolExecutorProtocol(Protocol):
+    def submit(self, fn: Callable, *args, **kwargs) -> Future:
+        ...
+
+    def set_size(self, max_workers: int) -> None:
+        """Set the size of the thread pool."""
+
+    def get_size(self) -> int:
+        """Return the current size of the thread pool."""
+
+
 class LazyThreadPool:
     """
     Lazily initialized thread pool.
@@ -23,6 +40,10 @@ class LazyThreadPool:
     _THREAD_POOL_FACTORY = ThreadPoolExecutor
 
     def __init__(self, max_workers: int | None = None, **kwargs):
+        if max_workers is None:
+            max_workers = min(
+                32, (os.cpu_count() or 1) + 4
+            )  # same default as in ThreadPoolExecutor
         self._max_workers = max_workers
         self._thread_pool: ThreadPoolExecutor | None = None
         super().__init__(**kwargs)
@@ -49,7 +70,8 @@ def set_size(self, max_workers: int) -> None:
             old_thread_pool.shutdown(wait=True)
         self._max_workers = max_workers
 
-    def get_size(self) -> int | None:
+    def get_size(self) -> int:
+        """Return the current size of the thread pool."""
         return self._max_workers
 
 
@@ -62,7 +84,7 @@ class ThreadPoolMixin(metaclass=B2TraceMetaAbstract):
 
     def __init__(
         self,
-        thread_pool: ThreadPoolExecutor | None = None,
+        thread_pool: DynamicThreadPoolExecutorProtocol | None = None,
         max_workers: int | None = None,
         **kwargs,
     ):
@@ -88,5 +110,5 @@ def set_thread_pool_size(self, max_workers: int) -> None:
         """
         return self._thread_pool.set_size(max_workers)
 
-    def get_thread_pool_size(self) -> int | None:
+    def get_thread_pool_size(self) -> int:
         return self._thread_pool.get_size()
diff --git a/test/unit/utils/test_thread_pool.py b/test/unit/utils/test_thread_pool.py
new file mode 100644
index 000000000..5742d3e35
--- /dev/null
+++ b/test/unit/utils/test_thread_pool.py
@@ -0,0 +1,43 @@
+######################################################################
+#
+# File: test/unit/utils/test_thread_pool.py
+#
+# Copyright 2024 Backblaze Inc. All Rights Reserved.
+#
+# License https://www.backblaze.com/using_b2_code.html
+#
+######################################################################
+from concurrent.futures import Future
+
+import pytest
+
+from b2sdk.utils.thread_pool import LazyThreadPool
+
+
+class TestLazyThreadPool:
+    @pytest.fixture
+    def thread_pool(self):
+        return LazyThreadPool()
+
+    def test_submit(self, thread_pool):
+
+        future = thread_pool.submit(sum, (1, 2))
+        assert isinstance(future, Future)
+        assert future.result() == 3
+
+    def test_set_size(self, thread_pool):
+        thread_pool.set_size(10)
+        assert thread_pool.get_size() == 10
+
+    def test_get_size(self, thread_pool):
+        assert thread_pool.get_size() > 0
+
+    def test_set_size__after_submit(self, thread_pool):
+        future = thread_pool.submit(sum, (1, 2))
+
+        thread_pool.set_size(7)
+        assert thread_pool.get_size() == 7
+
+        assert future.result() == 3
+
+        assert thread_pool.submit(sum, (1,)).result() == 1
diff --git a/test/unit/v_all/test_transfer.py b/test/unit/v_all/test_transfer.py
index 1492cac08..584bf6d68 100644
--- a/test/unit/v_all/test_transfer.py
+++ b/test/unit/v_all/test_transfer.py
@@ -19,16 +19,18 @@
 class TestDownloadManager(TestBase):
     def test_set_thread_pool_size(self) -> None:
         download_manager = DownloadManager(services=Mock())
-        assert download_manager.get_thread_pool_size() is None
-        download_manager.set_thread_pool_size(21)
-        assert download_manager._thread_pool._max_workers == 21
-        assert download_manager.get_thread_pool_size() == 21
+        assert download_manager.get_thread_pool_size() > 0
+
+        pool_size = 21
+        download_manager.set_thread_pool_size(pool_size)
+        assert download_manager.get_thread_pool_size() == pool_size
 
 
 class TestUploadManager(TestBase):
     def test_set_thread_pool_size(self) -> None:
         upload_manager = UploadManager(services=Mock())
-        assert upload_manager.get_thread_pool_size() is None
-        upload_manager.set_thread_pool_size(37)
-        assert upload_manager._thread_pool._max_workers == 37
-        assert upload_manager.get_thread_pool_size() == 37
+        assert upload_manager.get_thread_pool_size() > 0
+
+        pool_size = 37
+        upload_manager.set_thread_pool_size(pool_size)
+        assert upload_manager.get_thread_pool_size() == pool_size