From 59646ee1afbc14356e854cfe6c82bfe03f5c9f3a Mon Sep 17 00:00:00 2001 From: erezrokah Date: Tue, 20 Oct 2026 15:14:05 +0300 Subject: [PATCH 1/6] 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 221a722894d52bfd276ad520e2597695325bc545 Mon Sep 17 00:00:00 2001 From: Kemal Hadimli Date: Tue, 20 Oct 2026 15:15:38 +0100 Subject: [PATCH 2/6] fix: Send checkpoints flushed while a resource is resolved after it and its relations --- cloudquery/sdk/internal/checkpoint.py | 45 ++++++- cloudquery/sdk/scheduler/scheduler.py | 63 +++++----- tests/scheduler/checkpoints.py | 161 ++++++++++++++++++++++++++ 3 files changed, 234 insertions(+), 35 deletions(-) diff --git a/cloudquery/sdk/internal/checkpoint.py b/cloudquery/sdk/internal/checkpoint.py index 394d6a3..9f9b910 100644 --- a/cloudquery/sdk/internal/checkpoint.py +++ b/cloudquery/sdk/internal/checkpoint.py @@ -39,6 +39,10 @@ def __init__(self, parent: Optional[Scope], sync: Optional[object]): self.resolving = True # checkpoints waiting for this scope to finish self.waiters: List[_Pending] = [] + # set while a resource is resolved, until it is on the queue and its relations are submitted + self.resource_in_flight = False + # checkpoints waiting for that resource and its relations + self.after_resource: List[_Pending] = [] def finished(self) -> bool: return not self.resolving and not self.open_children @@ -112,6 +116,29 @@ def current_scope(self) -> Optional[Scope]: """Returns the scope of the table resolver running in this context, or None outside table resolvers.""" return self._scope.get() + @contextmanager + def resource(self, scope: Scope) -> Iterator[None]: + """ + Marks a resource of scope as in flight until it is on the queue and its relations are + submitted, so that checkpoints flushed meanwhile, e.g. from its resource hooks, follow + the resource and its relations. + """ + with self._lock: + scope.resource_in_flight = True + try: + yield + finally: + with self._lock: + scope.resource_in_flight = False + ready: List[_Pending] = [] + for pending in scope.after_resource: + pending.blockers -= 1 + self._wait_for_children(pending, scope) + if pending.blockers == 0: + ready.append(pending) + scope.after_resource = [] + self._send(ready) + def resolver_finished(self, scope: Scope) -> None: """ Marks the resolver of scope as returned, once every resource it produced is on the queue. @@ -167,11 +194,11 @@ def _schedule( ) -> None: pending = _Pending(msg, keys, self._sync) for scope in {s for s in scopes if s is not None}: - # the scope's own resources are on the queue already: resolvers put each - # resource on it before they resume, so only its relations can be in flight - for child in scope.open_children: - child.waiters.append(pending) + if scope.resource_in_flight: + scope.after_resource.append(pending) pending.blockers += 1 + else: + self._wait_for_children(pending, scope) earlier = set() for key in keys: previous = self._unsent.get(key) @@ -183,6 +210,16 @@ def _schedule( if pending.blockers == 0: self._send([pending]) + @staticmethod + def _wait_for_children(pending: _Pending, scope: Scope) -> None: + """ + Makes pending wait for the relations of scope's resources. Its resources are on the queue + already, as none is in flight. Called with the lock held. + """ + for child in scope.open_children: + child.waiters.append(pending) + pending.blockers += 1 + def _send(self, ready: List[_Pending]) -> None: """Puts ready checkpoints on the queue, then the ones that only waited for them. Called with the lock held.""" while ready: diff --git a/cloudquery/sdk/scheduler/scheduler.py b/cloudquery/sdk/scheduler/scheduler.py index f9f5cbb..bad46c8 100644 --- a/cloudquery/sdk/scheduler/scheduler.py +++ b/cloudquery/sdk/scheduler/scheduler.py @@ -134,37 +134,38 @@ def _resolve_table( depth=depth, ) for item in resolver.resolve(client, parent_item): - try: - resource = self.resolve_resource( - resolver, client, parent_item, item - ) - except Exception as e: - self._logger.error( - "failed to resolve resource", - client_id=client.id(), - table=resolver.table.name, - depth=depth, - exc_info=e, - ) - self._logger.debug( - "details about resource that failed to resolve", - client_id=client.id(), - table=resolver.table.name, - resource=item, - ) - continue - res.put(SyncInsertMessage(resource.to_arrow_record())) - for child_resolvers in resolver.child_resolvers: - res.put(TableResolverStarted()) - self._pools[depth + 1].submit( - self.resolve_table, - child_resolvers, - depth + 1, - client, - resource, - res, - self.checkpoints.new_scope(scope), - ) + with self.checkpoints.resource(scope): + try: + resource = self.resolve_resource( + resolver, client, parent_item, item + ) + except Exception as e: + self._logger.error( + "failed to resolve resource", + client_id=client.id(), + table=resolver.table.name, + depth=depth, + exc_info=e, + ) + self._logger.debug( + "details about resource that failed to resolve", + client_id=client.id(), + table=resolver.table.name, + resource=item, + ) + continue + res.put(SyncInsertMessage(resource.to_arrow_record())) + for child_resolvers in resolver.child_resolvers: + res.put(TableResolverStarted()) + self._pools[depth + 1].submit( + self.resolve_table, + child_resolvers, + depth + 1, + client, + resource, + res, + self.checkpoints.new_scope(scope), + ) total_resources += 1 if depth == 0: self._logger.info( diff --git a/tests/scheduler/checkpoints.py b/tests/scheduler/checkpoints.py index 61fc182..3ae279b 100644 --- a/tests/scheduler/checkpoints.py +++ b/tests/scheduler/checkpoints.py @@ -395,3 +395,164 @@ def test_resolvers_of_an_ended_sync_do_not_checkpoint_into_the_next_one(): ] assert cursors assert all(cursor.startswith("next-") for cursor in cursors) + + +class HookCursorResolver(TableResolver): + """Checkpoints its cursor from a resource hook, before the resource is on the queue.""" + + def __init__(self, state: StateClient, hook: str, parents: int): + super().__init__( + Table("parent", [Column("id", pa.int64())]), + child_resolvers=[ChildResolver(2)], + ) + self._state = state + self._hook = hook + self._parents = parents + + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: + for parent_id in range(self._parents): + yield {"id": parent_id} + + def pre_resource_resolve(self, client, resource): + if self._hook == "pre": + self._checkpoint(resource) + + def post_resource_resolve(self, client, resource): + if self._hook == "post": + self._checkpoint(resource) + + def _checkpoint(self, resource): + self._state.set_key("hook-cursor", str(resource.item["id"])) + self._state.flush() + + +@pytest.mark.parametrize("hook", ["pre", "post"]) +def test_checkpoints_flushed_by_resource_hooks_follow_the_resource_and_its_relations( + hook, +): + parents = 10 + scheduler = Scheduler(4) + state = MemoryStateClient(checkpoints=scheduler.checkpoints) + try: + messages = list( + scheduler.sync(TestClient(), [HookCursorResolver(state, hook, parents)]) + ) + finally: + scheduler.shutdown() + + seen = set() + cursors = [] + for msg in messages: + if isinstance(msg, SyncInsertMessage): + table = msg.record.schema.metadata[METADATA_TABLE_NAME] + row = msg.record.to_pylist()[0] + if table == b"parent": + seen.add((row["id"],)) + else: + seen.add((row["parent_id"], row["child_id"])) + elif isinstance(msg, SyncCheckpointMessage): + cursor = int(checkpoint_rows(msg)["hook-cursor"]) + cursors.append(cursor) + for parent_id in range(cursor + 1): + assert (parent_id,) in seen, f"cursor {cursor} sent before its parent" + for child_id in range(2): + assert (parent_id, child_id) in seen, ( + f"cursor {cursor} sent before child {child_id} of parent " + f"{parent_id}" + ) + assert cursors and cursors[-1] == parents - 1 + + +class GrandchildResolver(TableResolver): + """Yields slowly, so it still runs after the child resolver that spawned it returned.""" + + def __init__(self): + super().__init__( + Table( + "grandchild", + [ + Column("parent_id", pa.int64()), + Column("child_id", pa.int64()), + Column("grandchild_id", pa.int64()), + ], + ) + ) + + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: + for grandchild_id in range(3): + time.sleep(random.uniform(0.005, 0.02)) + yield {**parent_resource.item, "grandchild_id": grandchild_id} + + +class FastChildResolver(TableResolver): + def __init__(self): + super().__init__( + Table( + "fast_child", + [Column("parent_id", pa.int64()), Column("child_id", pa.int64())], + ), + child_resolvers=[GrandchildResolver()], + ) + + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: + for child_id in range(2): + yield {"parent_id": parent_resource.item["id"], "child_id": child_id} + + +class GrandparentResolver(TableResolver): + def __init__(self, state: StateClient, parents: int, flush_each: bool): + super().__init__( + Table("grandparent", [Column("id", pa.int64())]), + child_resolvers=[FastChildResolver()], + ) + self._state = state + self._parents = parents + self._flush_each = flush_each + + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: + for parent_id in range(self._parents): + yield {"id": parent_id} + if self._flush_each: + self._state.set_key("grandparent-cursor", str(parent_id)) + self._state.flush() + if not self._flush_each: + # the child resolvers return meanwhile, while their grandchildren still run + time.sleep(0.01) + self._state.set_key("grandparent-cursor", str(self._parents - 1)) + self._state.flush() + + +@pytest.mark.parametrize("flush_each", [True, False]) +def test_parent_checkpoints_follow_the_relations_of_their_relations(flush_each): + parents = 10 + scheduler = Scheduler(8) + state = MemoryStateClient(checkpoints=scheduler.checkpoints) + try: + messages = list( + scheduler.sync( + TestClient(), [GrandparentResolver(state, parents, flush_each)] + ) + ) + finally: + scheduler.shutdown() + + seen = set() + cursors = [] + for msg in messages: + if ( + isinstance(msg, SyncInsertMessage) + and msg.record.schema.metadata[METADATA_TABLE_NAME] == b"grandchild" + ): + row = msg.record.to_pylist()[0] + seen.add((row["parent_id"], row["child_id"], row["grandchild_id"])) + elif isinstance(msg, SyncCheckpointMessage): + cursor = int(checkpoint_rows(msg)["grandparent-cursor"]) + cursors.append(cursor) + for parent_id in range(cursor + 1): + for child_id in range(2): + for grandchild_id in range(3): + assert (parent_id, child_id, grandchild_id) in seen, ( + f"cursor {cursor} sent before grandchild " + f"{(parent_id, child_id, grandchild_id)}" + ) + assert cursors and cursors[-1] == parents - 1 From c278ea47d9f6f4990c44d07a72c40de0e3a2b718 Mon Sep 17 00:00:00 2001 From: Kemal Hadimli Date: Tue, 20 Oct 2026 15:16:43 +0100 Subject: [PATCH 3/6] fix: Raise SyncNotRunningError for checkpoints outside their sync, and skip unchanged state values --- cloudquery/sdk/internal/checkpoint.py | 14 +++++++++--- cloudquery/sdk/stateclient/__init__.py | 1 + cloudquery/sdk/stateclient/stateclient.py | 9 ++++++++ tests/scheduler/checkpoints.py | 28 ++++++++++++++++++++++- 4 files changed, 48 insertions(+), 4 deletions(-) diff --git a/cloudquery/sdk/internal/checkpoint.py b/cloudquery/sdk/internal/checkpoint.py index 9f9b910..11d3cf5 100644 --- a/cloudquery/sdk/internal/checkpoint.py +++ b/cloudquery/sdk/internal/checkpoint.py @@ -22,6 +22,14 @@ from cloudquery.sdk.message import SyncCheckpointMessage +class SyncNotRunningError(RuntimeError): + """ + Raised when state checkpoints are flushed while their sync is not running, so they cannot be + sent: before or after Scheduler.sync, or from a resolver of an ended sync. Flush state that + summarizes the whole sync in the scheduler's post-sync hook instead. + """ + + class Scope: """ A table resolver call and the relations its resources spawn. @@ -170,17 +178,17 @@ def schedule( wait for nothing. Schedules all of them or, if it raises, none of them. Raises: - RuntimeError: if no sync is running, as the checkpoints could not be sent. + SyncNotRunningError: if the sync of the checkpoints is not running. """ with self._lock: if self._res is None: - raise RuntimeError( + raise SyncNotRunningError( "state checkpoints can only be flushed while the scheduler syncs, " "e.g. from a table resolver or the post-sync hook" ) for scopes, _, _ in checkpoints: if any(s is not None and s.sync is not self._sync for s in scopes): - raise RuntimeError( + raise SyncNotRunningError( "state checkpoints of a table resolver can only be flushed during its own sync" ) for scopes, keys, msg in checkpoints: diff --git a/cloudquery/sdk/stateclient/__init__.py b/cloudquery/sdk/stateclient/__init__.py index ddff7be..319ee65 100644 --- a/cloudquery/sdk/stateclient/__init__.py +++ b/cloudquery/sdk/stateclient/__init__.py @@ -1 +1,2 @@ +from cloudquery.sdk.internal.checkpoint import SyncNotRunningError from .keys import state_key diff --git a/cloudquery/sdk/stateclient/stateclient.py b/cloudquery/sdk/stateclient/stateclient.py index dfa245a..c9b1109 100644 --- a/cloudquery/sdk/stateclient/stateclient.py +++ b/cloudquery/sdk/stateclient/stateclient.py @@ -87,6 +87,11 @@ def get_key(self, key: str) -> Optional[str]: def set_key(self, key: str, value: str) -> None: scope = self.checkpoints.current_scope() if self.checkpoints else None with self._lock: + if self.mem.get(key) == value: + # a pending change of the key summarizes this caller's resources too + if key in self.changes: + self.changes[key].add(scope) + return self.mem[key] = value self.changes.setdefault(key, set()).add(scope) @@ -95,6 +100,10 @@ def flush(self): Persists the keys changed since the last flush: writes them to the backend, or, with checkpoints, sends them on the sync stream after the resources they cover. If that fails, the keys stay changed, so a later flush retries them. + + Raises: + SyncNotRunningError: with checkpoints, if the sync is not running, e.g. once + Scheduler.sync returned. Flush in the scheduler's post-sync hook instead. """ with self._flush_lock: with self._lock: diff --git a/tests/scheduler/checkpoints.py b/tests/scheduler/checkpoints.py index 3ae279b..14706cf 100644 --- a/tests/scheduler/checkpoints.py +++ b/tests/scheduler/checkpoints.py @@ -18,6 +18,7 @@ from cloudquery.sdk.schema import Column from cloudquery.sdk.schema.arrow import METADATA_TABLE_NAME from cloudquery.sdk.schema.table import Table +from cloudquery.sdk.stateclient import SyncNotRunningError from cloudquery.sdk.stateclient.stateclient import ( ConnectedStateClient, StateClient, @@ -223,6 +224,31 @@ def test_flush_writes_to_the_backend_without_checkpoints(): assert state.writes[-1] == [("a", "3")] +def test_setting_an_unchanged_value_is_not_a_change(): + state = MemoryStateClient(checkpoints=None) + state.backend["a"] = "1" + state.read_all_state() + state.set_key("a", "1") + state.flush() + assert state.writes == [] + + state.set_key("a", "2") + state.set_key("a", "2") + state.flush() + assert state.writes == [[("a", "2")]] + + +def test_setting_an_unchanged_value_adds_the_caller_to_the_pending_change(): + scheduler = Scheduler(1) + state = MemoryStateClient(checkpoints=scheduler.checkpoints) + state.set_key("a", "1") + scope = scheduler.checkpoints.new_scope(None) + with scheduler.checkpoints.entered(scope): + state.set_key("a", "1") + scheduler.shutdown() + assert state.changes["a"] == {None, scope} + + def test_failed_flush_keeps_the_changes(): state = MemoryStateClient(checkpoints=None, fail_writes=True) state.set_key("a", "1") @@ -238,7 +264,7 @@ def test_checkpoint_flush_outside_a_sync_fails_and_keeps_the_changes(): scheduler = Scheduler(1) state = MemoryStateClient(checkpoints=scheduler.checkpoints) state.set_key("a", "1") - with pytest.raises(RuntimeError): + with pytest.raises(SyncNotRunningError): state.flush() assert "a" in state.changes assert state.writes == [] From 93c94a145687cd72937e3adedfa9adb7d9a45492 Mon Sep 17 00:00:00 2001 From: erezrokah Date: Tue, 20 Oct 2026 17:29:43 +0300 Subject: [PATCH 4/6] 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 From 70df295de0b180d34d1969688b06e49bdf621a68 Mon Sep 17 00:00:00 2001 From: Kemal Hadimli Date: Tue, 20 Oct 2026 15:45:29 +0100 Subject: [PATCH 5/6] refactor: Keep the pool's max_workers instead of reading ThreadPoolExecutor's private copy --- cloudquery/sdk/scheduler/scheduler.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/cloudquery/sdk/scheduler/scheduler.py b/cloudquery/sdk/scheduler/scheduler.py index d1ff1a8..b310adf 100644 --- a/cloudquery/sdk/scheduler/scheduler.py +++ b/cloudquery/sdk/scheduler/scheduler.py @@ -20,10 +20,12 @@ class ThreadPoolExecutorWithQueueSizeLimit(futures.ThreadPoolExecutor): - def __init__(self, maxsize, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, maxsize: int, max_workers: int, *args, **kwargs): + super().__init__(max_workers, *args, **kwargs) self.maxsize = maxsize - self._slots = threading.BoundedSemaphore(maxsize + self._max_workers) + # submit() puts on the work queue while holding a lock that every pool shares, so a + # full bounded queue would block the other pools' workers too: limit it before submit + self._slots = threading.BoundedSemaphore(maxsize + max_workers) def submit(self, fn, /, *args, **kwargs): self._slots.acquire() From 2346dea631998d9c33fd61b4c7de819037866b7b Mon Sep 17 00:00:00 2001 From: Kemal Hadimli Date: Tue, 20 Oct 2026 15:54:52 +0100 Subject: [PATCH 6/6] fix: Keep checkpoint order for unchanged state values --- cloudquery/sdk/stateclient/stateclient.py | 7 +++-- tests/scheduler/checkpoints.py | 32 +++++++++++++++++++++++ 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/cloudquery/sdk/stateclient/stateclient.py b/cloudquery/sdk/stateclient/stateclient.py index 502d4d1..a8550cd 100644 --- a/cloudquery/sdk/stateclient/stateclient.py +++ b/cloudquery/sdk/stateclient/stateclient.py @@ -91,10 +91,9 @@ def set_key(self, key: str, value: str) -> None: if self.checkpoints and self.checkpoints.superseded(scope): return with self._lock: - if self.mem.get(key) == value: - # a pending change of the key summarizes this caller's resources too - if key in self.changes: - self.changes[key].add(scope) + # with checkpoints, the value also covers the caller's resources, so later checkpoints + # of the key must still follow them even if it is unchanged + if self.checkpoints is None and self.mem.get(key) == value: return self.mem[key] = value self.changes.setdefault(key, set()).add(scope) diff --git a/tests/scheduler/checkpoints.py b/tests/scheduler/checkpoints.py index 24e542d..45ce726 100644 --- a/tests/scheduler/checkpoints.py +++ b/tests/scheduler/checkpoints.py @@ -1,4 +1,5 @@ import contextvars +import queue import random import threading import time @@ -249,6 +250,37 @@ def test_setting_an_unchanged_value_adds_the_caller_to_the_pending_change(): assert state.changes["a"] == {None, scope} +def test_unchanged_values_order_later_checkpoints_after_the_callers_resources(): + scheduler = Scheduler(1) + checkpoints = scheduler.checkpoints + res = queue.Queue() + checkpoints.start(res) + state = MemoryStateClient(checkpoints=checkpoints) + first, unfinished, last = (checkpoints.new_scope(None) for _ in range(3)) + child = checkpoints.new_scope(unfinished) + + def flush_in(scope, value): + with checkpoints.entered(scope): + state.set_key("cursor", value) + state.flush() + + def sent(): + return [checkpoint_rows(res.get_nowait())["cursor"] for _ in range(res.qsize())] + + try: + flush_in(first, "1") + assert sent() == ["1"] + # the same value, but it also covers the resources of a resolver whose relations still run + flush_in(unfinished, "1") + flush_in(last, "2") + assert sent() == [], "a later value must not overtake the relations it follows" + checkpoints.resolver_finished(child) + assert sent() == ["1", "2"] + finally: + checkpoints.stop() + scheduler.shutdown() + + def test_failed_flush_keeps_the_changes(): state = MemoryStateClient(checkpoints=None, fail_writes=True) state.set_key("a", "1")