diff --git a/cloudquery/sdk/scheduler/scheduler.py b/cloudquery/sdk/scheduler/scheduler.py index 6ea9abe..8c9c1ab 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 @@ -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( 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 new file mode 100644 index 0000000..94ea037 --- /dev/null +++ b/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