fix: Apply queue backpressure before submit instead of bounding the work queue

Author: erezrokah

cloudquery/sdk/scheduler/scheduler.py

@@ -1,4 +1,5 @@
 import queue
+import threading
 from concurrent import futures
 from typing import List, Generator, Any, Optional
 
@@ -20,7 +21,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:

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

@@ -1,61 +1,36 @@
-import gc
-import threading
-from typing import Any, Generator
+import subprocess
+import sys
 
-import pyarrow as pa
+import pytest
 
-from cloudquery.sdk.scheduler import Scheduler, TableResolver
+from cloudquery.sdk.scheduler import Scheduler
 from cloudquery.sdk.scheduler.scheduler import QUEUE_PER_WORKER
-from cloudquery.sdk.schema import Column
-from cloudquery.sdk.schema.table import Table
-
-
-class ParentResolver(TableResolver):
-    def __init__(self) -> None:
-        super().__init__(
-            Table("parent_table", [Column("id", pa.int64())]),
-            child_resolvers=[ChildResolver()],
-        )
-
-    def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
-        for i in range(3):
-            yield {"id": i}
-
-
-class ChildResolver(TableResolver):
-    def __init__(self) -> None:
-        super().__init__(Table("child_table", [Column("id", pa.int64())]))
-
-    def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
-        yield {"id": 1}
-
-
-class QueueSizeTestClient:
-    def id(self):
-        return "test_client"
 
 
 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._work_queue.maxsize == expected_size
+        assert pool.maxsize == expected_size
         expected_size = expected_size // 2 if expected_size > 1 else 1
     s.shutdown()
 
 
-def test_dropping_used_scheduler_does_not_block():
-    holder = {}
-    for _ in range(10):
-        holder["scheduler"] = Scheduler(10)
-        list(holder["scheduler"].sync(QueueSizeTestClient(), [ParentResolver()]))
-
-        def replace_scheduler():
-            holder["scheduler"] = Scheduler(10)
-            gc.collect()
-
-        thread = threading.Thread(target=replace_scheduler, daemon=True)
-        thread.start()
-        thread.join(timeout=5)
-        assert not thread.is_alive()
+@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