fix(checkpointer): serialise concurrent manual_checkpoint() flushes (#80) (#96)

Two asyncio tasks calling manual_checkpoint() against the same
BaseHeartbeatCheckPointer could both iterate the buffer snapshot and
double-write the same shard. On a last-writer-wins backend (Redis
GETSET), an older sequence could clobber a newer one. Without the lock
the second flusher could also KeyError on a shard the first had popped
mid-iteration.

Add asyncio.Lock around the manual_checkpoint() body so concurrent
callers serialise. Single-task callers see no behaviour change.

Document the full concurrency contract on the class so callers know
which methods are safe across tasks (only manual_checkpoint is
internally synchronised; checkpoint with auto_checkpoint=True,
allocate, deallocate, and close all need single-task drivers or
external coordination).

Hazard 3 from #80 (heartbeat dict-during-iteration) was closed
separately by #92 (issue #81 papercut bundle).

Author: hampsterx

kinesis/checkpointers.py

@@ -172,6 +172,38 @@ def is_allocated(self, shard_id):
 
 
 class BaseHeartbeatCheckPointer(BaseCheckPointer):
+    """Heartbeat-driven base for backend checkpointers (Redis, DynamoDB).
+
+    Concurrency contract:
+
+    * ``manual_checkpoint`` is the only method with explicit internal
+      serialisation. Concurrent callers from different tasks are guaranteed
+      not to double-flush the same buffered shard: the second caller waits
+      on ``_manual_flush_lock``, then iterates whatever the first left in
+      the buffer (failed shards retry, succeeded shards are gone). The lock
+      is held across backend writes, so a slow backend stalls concurrent
+      flushes.
+    * ``checkpoint`` with ``auto_checkpoint=False`` is a single dict
+      assignment to ``_manual_checkpoints`` and is safe to call from any
+      task. With ``auto_checkpoint=True`` it awaits ``_checkpoint`` directly
+      and offers no per-shard ordering guarantee against another concurrent
+      ``checkpoint`` for the same shard; on a last-writer-wins backend
+      (Redis ``GETSET``) two concurrent writers can race. The Consumer
+      drives this from a single task by design; callers building their own
+      driver must serialise per-shard themselves.
+    * ``allocate`` / ``deallocate`` perform backend writes and mutate
+      ``_items``. Concurrent calls for the *same* shard are not synchronised
+      and would race the backend lease/release; calls for *different* shards
+      are independent. The Consumer's shard manager invokes these from a
+      single task.
+    * ``heartbeat`` is owned by this class. The task is started in
+      ``__init__`` and torn down in ``close()``; do not invoke it directly.
+    * ``close`` should be called from a single task. Running it concurrently
+      with itself or with ``manual_checkpoint`` is undefined: ``close`` does
+      not take ``_manual_flush_lock`` and may deallocate a shard that an
+      in-flight flush is still writing.
+    """
+
     def __init__(
         self,
         name,
@@ -186,6 +218,7 @@ def __init__(
         self.heartbeat_frequency = heartbeat_frequency
         self.auto_checkpoint = auto_checkpoint
         self._manual_checkpoints: Dict[str, str] = {}
+        self._manual_flush_lock: asyncio.Lock = asyncio.Lock()
 
         self.heartbeat_task = asyncio.Task(self.heartbeat())
 
@@ -244,6 +277,12 @@ async def manual_checkpoint(self) -> None:
         so a concurrent ``checkpoint()`` that buffered a newer sequence for
         the same shard mid-flight is not silently dropped.
 
+        Concurrent callers serialise via ``_manual_flush_lock``: two tasks
+        calling ``manual_checkpoint()`` cannot both flush the same buffered
+        shard, which would otherwise risk an older sequence clobbering a
+        newer one on a last-writer-wins backend. See the class-level
+        thread-safety contract for the full picture.
+
         Raises:
             CheckpointFlushError: if one or more shards raised. The compound
                 exception carries the per-shard failures in ``.errors``.
@@ -252,18 +291,19 @@ async def manual_checkpoint(self) -> None:
                 callers should monitor that metric for alerting rather than
                 parsing the exception.
         """
-        errors: List[Tuple[str, BaseException]] = []
-        for shard_id in list(self._manual_checkpoints):
-            sequence = self._manual_checkpoints[shard_id]
-            try:
-                await self._checkpoint(shard_id, sequence)
-            except Exception as exc:
-                errors.append((shard_id, exc))
-                continue
-            if self._manual_checkpoints.get(shard_id) == sequence:
-                self._manual_checkpoints.pop(shard_id, None)
-        if errors:
-            raise CheckpointFlushError(errors) from errors[0][1]
+        async with self._manual_flush_lock:
+            errors: List[Tuple[str, BaseException]] = []
+            for shard_id in list(self._manual_checkpoints):
+                sequence = self._manual_checkpoints[shard_id]
+                try:
+                    await self._checkpoint(shard_id, sequence)
+                except Exception as exc:
+                    errors.append((shard_id, exc))
+                    continue
+                if self._manual_checkpoints.get(shard_id) == sequence:
+                    self._manual_checkpoints.pop(shard_id, None)
+            if errors:
+                raise CheckpointFlushError(errors) from errors[0][1]
 
 
 class MemoryCheckPointer(BaseCheckPointer):

tests/test_checkpoint_base_contract.py

@@ -186,6 +186,71 @@ async def inject_newer(shard_id: str, sequence: str) -> None:
         assert cp.writes == {"shard-0": "seq-1"}
 
 
+class TestManualCheckpointConcurrency:
+    """Issue #80 hazard 2: two tasks calling ``manual_checkpoint()`` must not
+    double-flush the same buffered shard. The base-class lock around the
+    flush body serialises them; only the first flusher writes any given
+    shard's buffered sequence, the second observes a drained buffer."""
+
+    @pytest.mark.asyncio
+    async def test_concurrent_flushes_serialise_per_shard(self, cp):
+        """Two ``manual_checkpoint()`` coroutines see each shard written
+        exactly once across the pair, in a deterministic order pinned by an
+        ``asyncio.Event`` injected mid-flush.
+
+        Without the lock both flushers iterate the same buffer snapshot and
+        each calls ``_checkpoint(shard, seq)`` once, producing two writes
+        per shard. With the lock the second flusher waits, sees the
+        shard already popped, and writes nothing.
+        """
+
+        await cp.checkpoint("shard-0", "seq-1")
+        await cp.checkpoint("shard-1", "seq-2")
+
+        write_log: list = []
+        gate = asyncio.Event()
+        first_in_flight = asyncio.Event()
+
+        async def gated_checkpoint(shard_id: str, sequence: str) -> None:
+            write_log.append((shard_id, sequence))
+            cp.writes[shard_id] = sequence
+            cp._items[shard_id] = sequence
+            if not first_in_flight.is_set():
+                first_in_flight.set()
+                await gate.wait()
+            cp._emit_checkpoint_success(shard_id)
+
+        cp._checkpoint = gated_checkpoint  # type: ignore[method-assign]
+
+        flush_a = asyncio.create_task(cp.manual_checkpoint())
+        await first_in_flight.wait()
+        flush_b = asyncio.create_task(cp.manual_checkpoint())
+        # Yield enough times for flush_b to reach the lock and block on it.
+        for _ in range(5):
+            await asyncio.sleep(0)
+        gate.set()
+        await asyncio.gather(flush_a, flush_b)
+
+        assert sorted(write_log) == [("shard-0", "seq-1"), ("shard-1", "seq-2")]
+        assert cp._manual_checkpoints == {}
+        assert cp.writes == {"shard-0": "seq-1", "shard-1": "seq-2"}
+
+    @pytest.mark.asyncio
+    async def test_second_flusher_sees_drained_buffer(self, cp):
+        """A second ``manual_checkpoint()`` arriving after the first drains
+        the buffer must return cleanly: empty buffer is a no-op, not an
+        error."""
+
+        await cp.checkpoint("shard-0", "seq-1")
+
+        await cp.manual_checkpoint()
+        # Second call against an already-drained buffer.
+        await cp.manual_checkpoint()
+
+        assert cp.writes == {"shard-0": "seq-1"}
+        assert cp._manual_checkpoints == {}
+
+
 class TestHeartbeatRobustness:
     """Issue #81: heartbeat iteration and close() must survive concurrent
     mutation / task death so shards don't leak."""