Skip to content

Commit bbd5b2e

Browse files
fix: Reject state changes from resolvers of an ended sync (#417)
* fix: Reject state changes from resolvers of an ended sync A resolver still running after its sync ended could change the state client of the next sync, e.g. one a plugin keeps on an attribute. Its key then failed every flush of that sync, as a failed flush keeps its changes, and the next sync read its value. set_key now raises before it records the change. * fix: Drop state changes of ended syncs instead of rejecting them --------- Co-authored-by: erezrokah <erezrokah@users.noreply.github.com>
1 parent 39aaf90 commit bbd5b2e

3 files changed

Lines changed: 134 additions & 3 deletions

File tree

‎cloudquery/sdk/internal/checkpoint.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,14 @@ def current_scope(self) -> Optional[Scope]:
112112
"""Returns the scope of the table resolver running in this context, or None outside table resolvers."""
113113
return self._scope.get()
114114

115+
def superseded(self, scope: Optional[Scope]) -> bool:
116+
"""Returns True if scope belongs to a sync that ended while another sync runs now."""
117+
return (
118+
scope is not None
119+
and self._sync is not None
120+
and scope.sync is not self._sync
121+
)
122+
115123
def resolver_finished(self, scope: Scope) -> None:
116124
"""
117125
Marks the resolver of scope as returned, once every resource it produced is on the queue.
@@ -152,7 +160,7 @@ def schedule(
152160
"e.g. from a table resolver or the post-sync hook"
153161
)
154162
for scopes, _, _ in checkpoints:
155-
if any(s is not None and s.sync is not self._sync for s in scopes):
163+
if any(self.superseded(s) for s in scopes):
156164
raise RuntimeError(
157165
"state checkpoints of a table resolver can only be flushed during its own sync"
158166
)

‎cloudquery/sdk/stateclient/stateclient.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,13 +79,17 @@ def __init__(
7979

8080
self.migrate_state_table()
8181
self.read_all_state()
82+
self._flushed = dict(self.mem)
8283

8384
def get_key(self, key: str) -> Optional[str]:
8485
with self._lock:
86+
self._forget_changes_of_ended_syncs(key)
8587
return self.mem.get(key)
8688

8789
def set_key(self, key: str, value: str) -> None:
8890
scope = self.checkpoints.current_scope() if self.checkpoints else None
91+
if self.checkpoints and self.checkpoints.superseded(scope):
92+
return
8993
with self._lock:
9094
self.mem[key] = value
9195
self.changes.setdefault(key, set()).add(scope)
@@ -95,9 +99,13 @@ def flush(self):
9599
Persists the keys changed since the last flush: writes them to the backend, or, with
96100
checkpoints, sends them on the sync stream after the resources they cover.
97101
If that fails, the keys stay changed, so a later flush retries them.
102+
Changes made by table resolvers of an ended sync are dropped once another sync runs,
103+
as the resources they cover may never have been committed.
98104
"""
99105
with self._flush_lock:
100106
with self._lock:
107+
for key in list(self.changes):
108+
self._forget_changes_of_ended_syncs(key)
101109
if not self.changes:
102110
return
103111
changes, self.changes = self.changes, {}
@@ -112,6 +120,21 @@ def flush(self):
112120
for key, scopes in changes.items():
113121
self.changes.setdefault(key, set()).update(scopes)
114122
raise
123+
with self._lock:
124+
self._flushed.update(values)
125+
126+
def _forget_changes_of_ended_syncs(self, key: str) -> None:
127+
if self.checkpoints is None or key not in self.changes:
128+
return
129+
scopes = {s for s in self.changes[key] if not self.checkpoints.superseded(s)}
130+
if scopes:
131+
self.changes[key] = scopes
132+
return
133+
del self.changes[key]
134+
if key in self._flushed:
135+
self.mem[key] = self._flushed[key]
136+
else:
137+
self.mem.pop(key, None)
115138

116139
def _send_checkpoints(
117140
self, changes: Dict[str, Set[Optional[Scope]]], values: Dict[str, str]

‎tests/scheduler/checkpoints.py‎

Lines changed: 102 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,8 @@
3030
class MemoryStateClient(StateClient):
3131
"""A state client whose backend is a dict."""
3232

33-
def __init__(self, checkpoints=None, fail_writes=False):
34-
self.backend: Dict[str, str] = {}
33+
def __init__(self, checkpoints=None, fail_writes=False, backend=None):
34+
self.backend: Dict[str, str] = dict(backend or {})
3535
self.writes: List[List[Tuple[str, str]]] = []
3636
self.fail_writes = fail_writes
3737
super().__init__(
@@ -395,3 +395,103 @@ def test_resolvers_of_an_ended_sync_do_not_checkpoint_into_the_next_one():
395395
]
396396
assert cursors
397397
assert all(cursor.startswith("next-") for cursor in cursors)
398+
399+
400+
def test_resolvers_of_an_ended_sync_cannot_change_the_state_of_the_next_one():
401+
plugin_attributes: Dict[str, StateClient] = {}
402+
next_started, leftover_done = threading.Event(), threading.Event()
403+
leftover_errors: List[Exception] = []
404+
405+
class LeftoverResolver(TableResolver):
406+
def __init__(self):
407+
super().__init__(Table("leftover", [Column("id", pa.int64())]))
408+
409+
def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
410+
yield {"id": 1}
411+
try:
412+
assert next_started.wait(5), "the next sync did not start"
413+
plugin_attributes["state"].set_key("leftover_cursor", "ended")
414+
except Exception as e:
415+
leftover_errors.append(e)
416+
finally:
417+
leftover_done.set()
418+
419+
class NextResolver(CursorResolver):
420+
def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
421+
next_started.set()
422+
assert leftover_done.wait(5), "the leftover resolver did not finish"
423+
yield from super().resolve(client, parent_resource)
424+
425+
scheduler = Scheduler(4)
426+
try:
427+
ended = scheduler.sync(TestClient(), [LeftoverResolver()])
428+
while not isinstance(next(ended), SyncInsertMessage):
429+
pass
430+
ended.close()
431+
432+
plugin_attributes["state"] = MemoryStateClient(
433+
checkpoints=scheduler.checkpoints
434+
)
435+
scheduler.set_post_sync_hook(plugin_attributes["state"].flush)
436+
messages = list(
437+
scheduler.sync(
438+
TestClient(), [NextResolver(plugin_attributes["state"], "next")]
439+
)
440+
)
441+
finally:
442+
scheduler.shutdown()
443+
assert leftover_errors == []
444+
assert plugin_attributes["state"].get_key("leftover_cursor") is None
445+
checkpoints = [
446+
checkpoint_rows(m) for m in messages if isinstance(m, SyncCheckpointMessage)
447+
]
448+
assert [c["cursor"] for c in checkpoints] == [f"next-{i}" for i in range(1, 30)]
449+
450+
451+
class SetKeyThenWait(TableResolver):
452+
def __init__(self, state: StateClient, release: threading.Event):
453+
super().__init__(Table("set_then_wait", [Column("id", pa.int64())]))
454+
self._state = state
455+
self._release = release
456+
457+
def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
458+
self._state.set_key("cursor", "ended")
459+
yield {"id": 1}
460+
self._release.wait(5)
461+
462+
463+
def test_a_state_client_reused_across_syncs_drops_unflushed_changes_of_an_ended_sync():
464+
release = threading.Event()
465+
scheduler = Scheduler(4)
466+
try:
467+
state = MemoryStateClient(
468+
checkpoints=scheduler.checkpoints, backend={"cursor": "flushed"}
469+
)
470+
ended = scheduler.sync(TestClient(), [SetKeyThenWait(state, release)])
471+
while not isinstance(next(ended), SyncInsertMessage):
472+
pass
473+
ended.close()
474+
release.set()
475+
476+
scheduler.set_post_sync_hook(state.flush)
477+
cursors_seen: List[Optional[str]] = []
478+
479+
class ReadCursor(TableResolver):
480+
def __init__(self):
481+
super().__init__(Table("read_cursor", [Column("id", pa.int64())]))
482+
483+
def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
484+
cursors_seen.append(state.get_key("cursor"))
485+
yield {"id": 1}
486+
state.set_key("other", "next")
487+
488+
messages = list(scheduler.sync(TestClient(), [ReadCursor()]))
489+
list(scheduler.sync(TestClient(), [ReadCursor()]))
490+
finally:
491+
release.set()
492+
scheduler.shutdown()
493+
assert cursors_seen == ["flushed", "flushed"]
494+
checkpoints = [
495+
checkpoint_rows(m) for m in messages if isinstance(m, SyncCheckpointMessage)
496+
]
497+
assert checkpoints == [{"other": "next"}]

0 commit comments

Comments
 (0)