fix: Use the computed queue size for child resolver pools

cloudquery/sdk/scheduler/scheduler.py

@@ -1,4 +1,5 @@
 import queue
+import threading
 from concurrent import futures
 from typing import List, Generator, Any, Optional
 
@@ -21,7 +22,18 @@
 class ThreadPoolExecutorWithQueueSizeLimit(futures.ThreadPoolExecutor):
     def __init__(self, maxsize, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._work_queue = queue.Queue(maxsize=maxsize)
+        self.maxsize = maxsize
+        self._slots = threading.BoundedSemaphore(maxsize + self._max_workers)
+
+    def submit(self, fn, /, *args, **kwargs):
+        self._slots.acquire()
+        try:
+            future = super().submit(fn, *args, **kwargs)
+        except BaseException:
+            self._slots.release()
+            raise
+        future.add_done_callback(lambda _: self._slots.release())
+        return future
 
 
 class TableResolverStarted:
@@ -61,7 +73,7 @@ def __init__(
         )
         self._pools: List[ThreadPoolExecutorWithQueueSizeLimit] = []
         current_depth_concurrency = concurrency
-        current_depth_queue_size = queue_size
+        current_depth_queue_size = self._queue_size
         for _ in range(max_depth + 1):
             self._pools.append(
                 ThreadPoolExecutorWithQueueSizeLimit(

tests/scheduler/queue_size_scenarios.py

@@ -0,0 +1,84 @@
+import gc
+import sys
+import threading
+from typing import Any, Generator, List
+
+import pyarrow as pa
+
+from cloudquery.sdk.scheduler import Scheduler, TableResolver
+from cloudquery.sdk.schema import Column
+from cloudquery.sdk.schema.table import Table
+
+
+class ScenarioClient:
+    def __init__(self, index: int = 0) -> None:
+        self.index = index
+
+    def id(self):
+        return f"client_{self.index}"
+
+
+class ScenarioResolver(TableResolver):
+    def __init__(
+        self,
+        name: str,
+        items: int,
+        child_resolvers: List[TableResolver] = None,
+        clients: int = 1,
+    ) -> None:
+        super().__init__(
+            Table(name, [Column("id", pa.int64())]),
+            child_resolvers=child_resolvers or [],
+        )
+        self.items = items
+        self.clients = clients
+
+    def multiplex(self, client):
+        return [ScenarioClient(i) for i in range(self.clients)]
+
+    def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
+        for i in range(self.items):
+            yield {"id": i}
+
+
+def parent_with_child(clients: int = 1) -> ScenarioResolver:
+    return ScenarioResolver(
+        "parent", 3, [ScenarioResolver("child", 1)], clients=clients
+    )
+
+
+def sync_all(scheduler: Scheduler, resolver: TableResolver) -> None:
+    list(scheduler.sync(ScenarioClient(), [resolver]))
+
+
+def drop_used_scheduler(queue_size: int = 0) -> None:
+    holder = {}
+    for _ in range(10):
+        holder["scheduler"] = Scheduler(10, queue_size=queue_size)
+        sync_all(holder["scheduler"], parent_with_child())
+
+        def replace_scheduler():
+            holder["scheduler"] = Scheduler(10, queue_size=queue_size)
+            gc.collect()
+
+        thread = threading.Thread(target=replace_scheduler)
+        thread.start()
+        thread.join()
+
+
+def drop_used_scheduler_with_small_queue() -> None:
+    drop_used_scheduler(queue_size=1)
+
+
+def more_clients_than_top_level_queue() -> None:
+    sync_all(Scheduler(2), parent_with_child(clients=1000))
+
+
+def grandchildren_with_small_queue() -> None:
+    grandchild = ScenarioResolver("grandchild", 3)
+    child = ScenarioResolver("child", 3, [grandchild])
+    sync_all(Scheduler(4, queue_size=1), ScenarioResolver("parent", 10, [child]))
+
+
+if __name__ == "__main__":
+    globals()[sys.argv[1]]()

tests/scheduler/scheduler_queue_size.py

@@ -0,0 +1,36 @@
+import subprocess
+import sys
+
+import pytest
+
+from cloudquery.sdk.scheduler import Scheduler
+from cloudquery.sdk.scheduler.scheduler import QUEUE_PER_WORKER
+
+
+def test_default_queue_size_is_used_for_all_depths():
+    concurrency = 10
+    s = Scheduler(concurrency, max_depth=3)
+    expected_size = concurrency * QUEUE_PER_WORKER
+    for pool in s._pools:
+        assert pool.maxsize == expected_size
+        expected_size = expected_size // 2 if expected_size > 1 else 1
+    s.shutdown()
+
+
+@pytest.mark.parametrize(
+    "scenario",
+    [
+        "drop_used_scheduler",
+        "drop_used_scheduler_with_small_queue",
+        "more_clients_than_top_level_queue",
+        "grandchildren_with_small_queue",
+    ],
+)
+def test_scenario_does_not_hang(scenario):
+    result = subprocess.run(
+        [sys.executable, "-m", "tests.scheduler.queue_size_scenarios", scenario],
+        capture_output=True,
+        text=True,
+        timeout=60,
+    )
+    assert result.returncode == 0, result.stderr