From 59646ee1afbc14356e854cfe6c82bfe03f5c9f3a Mon Sep 17 00:00:00 2001 From: erezrokah Date: Tue, 20 Oct 2026 15:14:05 +0300 Subject: [PATCH 1/2] fix: Use the computed queue size for child resolver pools --- cloudquery/sdk/scheduler/scheduler.py | 2 +- tests/scheduler/scheduler_queue_size.py | 61 +++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) create mode 100644 tests/scheduler/scheduler_queue_size.py diff --git a/cloudquery/sdk/scheduler/scheduler.py b/cloudquery/sdk/scheduler/scheduler.py index bc00b99..64e34ed 100644 --- a/cloudquery/sdk/scheduler/scheduler.py +++ b/cloudquery/sdk/scheduler/scheduler.py @@ -53,7 +53,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( diff --git a/tests/scheduler/scheduler_queue_size.py b/tests/scheduler/scheduler_queue_size.py new file mode 100644 index 0000000..14e9a39 --- /dev/null +++ b/tests/scheduler/scheduler_queue_size.py @@ -0,0 +1,61 @@ +import gc +import threading +from typing import Any, Generator + +import pyarrow as pa + +from cloudquery.sdk.scheduler import Scheduler, TableResolver +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 + 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() From 93c94a145687cd72937e3adedfa9adb7d9a45492 Mon Sep 17 00:00:00 2001 From: erezrokah Date: Tue, 20 Oct 2026 17:29:43 +0300 Subject: [PATCH 2/2] fix: Apply queue backpressure before submit instead of bounding the work queue --- cloudquery/sdk/scheduler/scheduler.py | 14 ++++- tests/scheduler/queue_size_scenarios.py | 84 +++++++++++++++++++++++++ tests/scheduler/scheduler_queue_size.py | 69 +++++++------------- 3 files changed, 119 insertions(+), 48 deletions(-) create mode 100644 tests/scheduler/queue_size_scenarios.py diff --git a/cloudquery/sdk/scheduler/scheduler.py b/cloudquery/sdk/scheduler/scheduler.py index 64e34ed..23e4255 100644 --- a/cloudquery/sdk/scheduler/scheduler.py +++ b/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: diff --git a/tests/scheduler/queue_size_scenarios.py b/tests/scheduler/queue_size_scenarios.py new file mode 100644 index 0000000..a390a72 --- /dev/null +++ b/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]]() diff --git a/tests/scheduler/scheduler_queue_size.py b/tests/scheduler/scheduler_queue_size.py index 14e9a39..94ea037 100644 --- a/tests/scheduler/scheduler_queue_size.py +++ b/tests/scheduler/scheduler_queue_size.py @@ -1,38 +1,10 @@ -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(): @@ -40,22 +12,25 @@ def test_default_queue_size_is_used_for_all_depths(): 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