fix: Send checkpoints from resource hooks after their resources

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:

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(

cloudquery/sdk/stateclient/__init__.py

@@ -1 +1,2 @@
+from cloudquery.sdk.internal.checkpoint import SyncNotRunningError
 from .keys import state_key

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: