diff --git a/cloudquery/sdk/internal/checkpoint.py b/cloudquery/sdk/internal/checkpoint.py index 3e0dbf2..850eaba 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. @@ -39,6 +47,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 +124,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 superseded(self, scope: Optional[Scope]) -> bool: """Returns True if scope belongs to a sync that ended while another sync runs now.""" return ( @@ -151,17 +186,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(self.superseded(s) 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: @@ -175,11 +210,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) @@ -191,6 +226,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 8c9c1ab..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() @@ -146,37 +148,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/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 af817be..a8550cd 100644 --- a/cloudquery/sdk/stateclient/stateclient.py +++ b/cloudquery/sdk/stateclient/stateclient.py @@ -91,6 +91,10 @@ def set_key(self, key: str, value: str) -> None: if self.checkpoints and self.checkpoints.superseded(scope): return with self._lock: + # 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) @@ -101,6 +105,10 @@ def flush(self): If that fails, the keys stay changed, so a later flush retries them. Changes made by table resolvers of an ended sync are dropped once another sync runs, as the resources they cover may never have been committed. + + 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 8cf7741..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 @@ -18,6 +19,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 +225,62 @@ 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_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") @@ -238,7 +296,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 == [] @@ -495,3 +553,164 @@ def resolve(self, client, parent_resource) -> Generator[Any, None, None]: checkpoint_rows(m) for m in messages if isinstance(m, SyncCheckpointMessage) ] assert checkpoints == [{"other": "next"}] + + +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