fix(insight): flush a record already taken for export

The condition on the refused drain's flush request was pending-only, and the
worker takes a record out of the pending map before it calls the exporter. A
hook re-entered from inside export() that reached an invocation end emitting
no record therefore found nothing pending, asked for no flush, and left the
snapshot it had just been handed in a buffering exporter until the
environment froze -- the loss the request exists to prevent, arriving by the
other door.

The condition now separates the two re-entry paths, because they need
opposite answers: a flush already in flight with nothing newly pending is the
one case that asks for nothing, which is what stops an exporter whose flush()
re-enters from flushing for as long as the environment lives.

The plugin contract also said a plugin may hold "per-execution" state in its
attributes. Per invocation is narrower, and the difference is the replay
model: an execution spans as many invocations as it waits, retries or
resumes, and the instance is dropped when each returns. The README and the
factory protocol now say per-invocation and name the two homes for anything
that has to outlive it -- the operation-map snapshot the invocation hooks
carry, or the factory.

Author: ParidelPooya

packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py

@@ -301,22 +301,27 @@ def _is_export_worker(self) -> bool:
             return self._worker is threading.current_thread()
 
     def _request_flush_for_pending_records(self) -> None:
-        """Ask the worker for a flush, but only if a record is waiting for one.
+        """Ask the worker for a flush, unless a flush already in flight covers it.
 
         Returns without waiting, so it is safe to call from the worker itself. The
         barrier is raised to the current schedule counter, which is what makes the
         flush cover a record queued moments ago rather than running before it.
 
-        A pending record is the condition, not a formality. This is called from a
-        drain refused on the worker thread, and one way to reach that is an
-        exporter whose ``flush()`` re-enters a plugin hook: the call then arrives
-        from inside a flush, and requesting the next one unconditionally would
-        produce a flush that re-enters, requests, and flushes again for as long as
-        the environment lives -- after the invocation has returned. Nothing is
-        pending in that case, so nothing is requested.
+        The condition separates the two ways a refused drain is reached, because
+        they need opposite answers. Re-entered from an exporter's ``export()``,
+        the record has already left ``_pending`` -- the worker takes it before it
+        calls the exporter -- so a pending-only test would skip the request and
+        leave that snapshot buffered until the environment froze. Re-entered from
+        an exporter's ``flush()``, a flush is in flight and covers what was
+        exported before it, so requesting another would produce a flush that
+        re-enters, requests, and flushes again for as long as the environment
+        lived -- after the invocation returned. A flush in flight with nothing
+        newly pending is therefore the one case that asks for nothing.
         """
         with self._condition:
-            if self._disabled or not self._pending:
+            if self._disabled:
+                return
+            if not self._pending and self._flush_in_flight is not None:
                 return
             self._flush_requested = True
             self._flush_barrier = max(self._flush_barrier, self._seq)

packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py

@@ -968,3 +968,50 @@ def test_a_drain_refused_from_inside_a_flush_does_not_re_arm_it() -> None:
     # assertion: an unconditional request never settles.
     assert not _wait_until(lambda: exporter.flushes > flushes_after_drain, timeout=1.0)
     assert _wait_until(lambda: not scheduler._worker_alive())
+
+
+class ReentrantDrainDuringExportExporter(CaptureExporter):
+    """Exporter whose export() drains without queueing anything new.
+
+    The worker takes a record out of the pending map before it calls the
+    exporter, so a hook re-entered from inside ``export()`` and reaching an
+    invocation end that emits no record finds nothing pending -- while the record
+    it was just handed is still only in the exporter's buffer.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.scheduler: _ArnScheduler | None = None
+        self.returned = threading.Event()
+        self._reentered = False
+
+    def export(self, record: dict[str, Any]) -> None:
+        super().export(record)
+        assert self.scheduler is not None
+        if self._reentered:
+            return
+        self._reentered = True
+        self.scheduler.drain(ARN_B)
+        self.returned.set()
+
+
+def test_a_drain_refused_from_inside_an_export_still_flushes() -> None:
+    """A refused drain from inside an export asks for a flush.
+
+    The record it must cover has already left the pending map, so a
+    pending-only condition would leave that snapshot in a buffering exporter
+    when the environment froze -- which is the loss the refusal path exists to
+    prevent, arriving by the other door.
+    """
+    exporter = ReentrantDrainDuringExportExporter()
+    scheduler = _ArnScheduler([exporter])
+    exporter.scheduler = scheduler
+
+    scheduler.schedule(ARN_A, _record("r1"))
+
+    assert exporter.returned.wait(timeout=10), "the refused drain must return"
+    assert _wait_until(lambda: ("flush", None) in exporter.calls), (
+        "the exported record must be flushed even though nothing was pending"
+    )
+    assert ("export", "r1") in exporter.calls
+    assert _wait_until(lambda: not scheduler._worker_alive())

packages/aws-durable-execution-sdk-python/README.md

@@ -47,7 +47,18 @@ A plugin is registered as a *factory*, not as an instance. A factory is an objec
 with a `create_plugin(info)` method taking the invocation's `InvocationStartInfo`
 and returning a `DurableInstrumentationPlugin`; the SDK calls that method once per
 invocation, so the instance it returns serves that one invocation only and can
-hold per-execution state in ordinary attributes.
+hold **per-invocation** state in ordinary attributes.
+
+Per-invocation is narrower than per-execution, and the difference matters. A
+durable execution spans as many invocations as it waits, retries or resumes, and
+the instance is dropped when each of those returns — so anything a plugin keeps in
+its attributes is gone by the next invocation of the same execution. State that
+has to survive that has two honest homes: rebuild it from the operation map the
+invocation hooks carry (`InvocationStartInfo.operations` is a full snapshot,
+which is how the bundled Insight plugin reports operations that completed in an
+earlier invocation), or put it on the factory, which outlives every invocation —
+keyed by execution ARN, and pruned by the owner, because the SDK will not tell the
+factory when an execution ends for good.
 
 A factory is an object with a method rather than a plain callable so the
 registration type can grow a second, optional member later -- a process-level

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py

@@ -505,9 +505,18 @@ def create_plugin(
         :class:`InvocationStartInfo` -- the same object the returned instance's
         ``on_invocation_start`` then receives -- and before any hook fires. The
         instance serves only that invocation and is dropped when it returns, so a
-        plugin can hold per-execution state in ordinary instance attributes
+        plugin can hold that invocation's state in ordinary instance attributes
         without keying it by execution ARN.
 
+        Per invocation is narrower than per execution. A durable execution spans
+        as many invocations as it waits, retries or resumes, so state a plugin
+        leaves in its attributes is gone by the next invocation of the same
+        execution. Anything that has to survive that is rebuilt from the operation
+        map the invocation hooks carry -- ``InvocationStartInfo.operations`` is a
+        full snapshot, including operations that completed in an earlier
+        invocation -- or kept on the factory, which outlives every invocation and
+        is therefore the caller's to key and to prune.
+
         ``info`` is positional-only, so an implementation may name the parameter
         whatever reads best; a named protocol parameter would pin that name for
         every implementation.