Merge pull request #41 from aerospike/dev

feat: [AIE-3] Use atomic operations instead of read-modify-write

Author: DKZed

.github/workflows/upgrade-lock.yml

@@ -56,10 +56,6 @@ jobs:
           add-paths: uv.lock
           commit-message: "chore: refresh uv.lock via weekly upgrade"
           title: "chore: weekly uv lockfile upgrade"
-          # `main` requires signed commits (see setup-github-protection.sh).
-          # `sign-commits: true` makes the action create the commit via the
-          # GitHub Contents API, which signs it with GitHub's web-flow GPG
-          # key so it shows as "Verified" and satisfies the rule.
           sign-commits: true
           body: |
             Automated weekly run of `uv lock --upgrade`.

packages/langgraph-checkpoint-aerospike/langgraph/checkpoint/aerospike/saver.py

@@ -3,11 +3,11 @@
 import asyncio
 import builtins
 import contextlib
-import json
 from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
 from datetime import datetime, timezone
 from typing import Any, cast
 
+from aerospike_helpers.operations import map_operations, operations
 from langchain_core.runnables import RunnableConfig
 from langgraph.checkpoint.base import (
     WRITES_IDX_MAP,
@@ -114,13 +114,19 @@ def _key_timeline(self, thread_id: str, checkpoint_ns: str):
         return (self.ns, self.set_meta, f"{thread_id}{SEP}{checkpoint_ns}{SEP}__timeline__")
 
     # ---------- aerospike io ----------
-    def _put(self, key, bins: dict[str, Any]) -> None:
-        policy: dict[str, Any] | None = None
+    def _ttl_policy(self) -> dict[str, Any] | None:
+        """Return ``{"ttl": seconds}`` for the configured TTL, or ``None``.
+
+        Passed as ``policy=`` to both ``client.put`` and ``client.operate``.
+        """
         minutes = self._ttl_minutes
-        if minutes is not None:
-            minutes = int(minutes)
-            if minutes > 0:
-                policy = {"ttl": minutes * 60}
+        if minutes is None:
+            return None
+        seconds = int(minutes) * 60
+        return {"ttl": seconds} if seconds > 0 else None
+
+    def _put(self, key, bins: dict[str, Any]) -> None:
+        policy = self._ttl_policy()
         try:
             if policy is not None:
                 self.client.put(key, bins, policy=policy)
@@ -130,39 +136,45 @@ def _put(self, key, bins: dict[str, Any]) -> None:
             raise RuntimeError(f"Aerospike put failed for {key}: {e}") from e
 
     def _get(self, key) -> tuple | None:
+        # `read_touch_ttl_percent=100` refreshes the TTL on every
+        # successful read, server-side, in the same round-trip.
+        policy: dict[str, Any] | None = None
+        if self._refresh_on_read and self._ttl_minutes is not None and self._ttl_minutes > 0:
+            policy = {"read_touch_ttl_percent": 100}
+
         try:
-            rec = self.client.get(key)
+            if policy is not None:
+                rec = self.client.get(key, policy=policy)
+            else:
+                rec = self.client.get(key)
         except aerospike.exception.RecordNotFound:
             return None
         except aerospike.exception.AerospikeError as e:
             raise RuntimeError(f"Aerospike get failed for {key}: {e}") from e
 
-        if self._refresh_on_read and self._ttl_minutes is not None and self._ttl_minutes > 0:
-            with contextlib.suppress(aerospike.exception.AerospikeError):
-                self.client.touch(key, int(self._ttl_minutes) * 60)
-
         return rec
 
     def _read_timeline_items(self, timeline_key) -> builtins.list[tuple[str, str]]:
-        """Return timeline entries as ``(iso_timestamp, checkpoint_id)`` pairs."""
+        """Return timeline entries as ``(iso_timestamp, checkpoint_id)`` pairs.
+
+        On-disk shape is a Map bin (``timeline``) keyed by
+        ``checkpoint_id`` with ISO-timestamp values. We sort by ``ts``
+        descending here so callers see reverse-chronological order.
+        """
         rec = self._get(timeline_key)
         if rec is None:
             return []
         bins = rec[2]
-        try:
-            items = json.loads(bins.get("items", "[]"))
-            cleaned: list[tuple[str, str]] = []
-            for it in items:
-                if (
-                    isinstance(it, list)
-                    and len(it) == 2
-                    and isinstance(it[0], str)
-                    and isinstance(it[1], str)
-                ):
-                    cleaned.append((it[0], it[1]))
-            return cleaned
-        except Exception:
+        timeline = bins.get("timeline") or {}
+        if not isinstance(timeline, dict):
             return []
+        pairs: list[tuple[str, str]] = [
+            (ts, cid)
+            for cid, ts in timeline.items()
+            if isinstance(ts, str) and isinstance(cid, str)
+        ]
+        pairs.sort(key=lambda p: p[0], reverse=True)
+        return pairs
 
     def _delete(self, key) -> None:
         try:
@@ -220,19 +232,21 @@ def put(
         )
 
         timeline_key = self._key_timeline(thread_id, checkpoint_ns)
-        items = self._read_timeline_items(timeline_key)
-
-        items = [(t, cid) for (t, cid) in items if cid != checkpoint_id]
-        items.insert(0, (ts, checkpoint_id))
-        if self.timeline_max is not None and len(items) > self.timeline_max:
-            items = items[: self.timeline_max]
-        self._put(
-            timeline_key,
-            {
-                "thread_id": thread_id,
-                "items": json.dumps(items),
-            },
-        )
+        # `map_put` upserts atomically: re-`put()`ing the same
+        # `checkpoint_id` overwrites in place, and concurrent `put()`s
+        # against the same thread/ns can't clobber each other's entries.
+        timeline_ops: list[dict[str, Any]] = [
+            operations.write("thread_id", thread_id),
+            map_operations.map_put("timeline", checkpoint_id, ts),
+        ]
+        timeline_policy = self._ttl_policy()
+        try:
+            if timeline_policy is not None:
+                self.client.operate(timeline_key, timeline_ops, policy=timeline_policy)
+            else:
+                self.client.operate(timeline_key, timeline_ops)
+        except aerospike.exception.AerospikeError as e:
+            raise RuntimeError(f"Aerospike operate failed for {timeline_key}: {e}") from e
 
         cfg_conf: dict[str, Any] = {**(config.get("configurable") or {})}
         cfg_conf.update(
@@ -252,30 +266,31 @@ def put_writes(
         task_id: str,
         task_path: str = "",
     ) -> None:
+        """Persist pending writes for a checkpoint.
 
+        Each write is stored inside a Map bin (``writes``) keyed by
+        ``f"{task_id}|{idx}"``, written via a single ``client.operate``
+        call. ``map_put`` is server-atomic, giving us upsert-on-retry
+        and tolerating concurrent callers against the same checkpoint.
+
+        The ``thread_id`` bin is rewritten on every call so that
+        ``delete_thread``'s secondary-index query keeps finding the
+        record; the value is the same every time for a given key.
+        """
         if not writes:
             return
 
         thread_id, checkpoint_ns, checkpoint_id = self._ids_from_config(config)
         if not checkpoint_id:
             return
 
-        key = self._key_writes(
-            thread_id, checkpoint_ns, checkpoint_id
-        )  # Do we need to put task id as key. (Reason: Record limit 8Mb)
-
-        existing_rec = self._get(key)
-        existing_items: list[dict[str, Any]] = []
-        if existing_rec is not None:
-            _, _, bins = existing_rec
-            existing_items = bins.get("writes") or []
-
+        key = self._key_writes(thread_id, checkpoint_ns, checkpoint_id)
         now_ts = _now_ns().isoformat()
 
+        ops: list[dict[str, Any]] = [operations.write("thread_id", thread_id)]
         for idx, (channel, value) in enumerate(writes):
             idx_val = WRITES_IDX_MAP.get(channel, idx)
             type_, serialized = self.serde.dumps_typed(value)
-
             new_item = {
                 "task_id": task_id,
                 "task_path": task_path,
@@ -285,24 +300,17 @@ def put_writes(
                 "value": serialized,
                 "ts": now_ts,
             }
-            replace_at: int | None = None
-            for i, item in enumerate(existing_items):
-                if item.get("task_id") == task_id and item.get("idx") == idx_val:
-                    replace_at = i
-                    break
-
-            if replace_at is not None:
-                existing_items[replace_at] = new_item
-            else:
-                existing_items.append(new_item)
+            map_key = f"{task_id}{SEP}{idx_val}"
+            ops.append(map_operations.map_put("writes", map_key, new_item))
 
-        self._put(
-            key,
-            {
-                "thread_id": thread_id,
-                "writes": existing_items,
-            },
-        )
+        policy = self._ttl_policy()
+        try:
+            if policy is not None:
+                self.client.operate(key, ops, policy=policy)
+            else:
+                self.client.operate(key, ops)
+        except aerospike.exception.AerospikeError as e:
+            raise RuntimeError(f"Aerospike operate failed for {key}: {e}") from e
 
     def get_tuple(
         self,
@@ -346,8 +354,11 @@ def get_tuple(
         wrec = self._get(self._key_writes(thread_id, checkpoint_ns, checkpoint_id))
         if wrec is not None:
             _, _, wbins = wrec
-            items = wbins.get("writes") or []
-            for item in items:
+            # `writes` is a Map bin (see `put_writes`); each value
+            # carries its own `task_id`, `channel`, and `idx`, so we
+            # don't depend on map iteration order.
+            writes_map = wbins.get("writes") or {}
+            for item in writes_map.values():
                 try:
                     task_id = item.get("task_id", "")
                     channel = item["channel"]

packages/langgraph-checkpoint-aerospike/tests/test_saver_concurrency.py

@@ -0,0 +1,121 @@
+"""Concurrency tests for ``AerospikeSaver``.
+
+The upstream conformance suite is single-threaded and the integration
+tests won't reliably surface contention bugs, so we cover three
+properties directly here:
+
+* ``put_writes`` from N concurrent callers against the same checkpoint
+  produces N distinct pending writes (no lost updates).
+* ``put_writes`` retries of the same ``(task_id, idx)`` upsert in place
+  rather than appending duplicates.
+* ``put()`` from N concurrent callers against the same thread/ns
+  produces N distinct timeline entries.
+"""
+
+from __future__ import annotations
+
+import threading
+from concurrent.futures import ThreadPoolExecutor
+
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import empty_checkpoint
+
+
+def _seed_checkpoint(saver, thread_id: str, checkpoint_ns: str = "") -> RunnableConfig:
+    """Write an empty checkpoint so subsequent ``put_writes`` has a target."""
+    base_config: RunnableConfig = {
+        "configurable": {
+            "thread_id": thread_id,
+            "checkpoint_ns": checkpoint_ns,
+        }
+    }
+    cp = empty_checkpoint()
+    metadata = {"source": "input", "step": 0, "writes": {}, "parents": {}}
+    new_config = saver.put(base_config, cp, metadata, {})  # type: ignore[arg-type]
+    return new_config
+
+
+def test_put_writes_concurrent_no_lost_updates(saver) -> None:
+    """N parallel ``put_writes`` calls must each be visible in pending_writes."""
+    cp_config = _seed_checkpoint(saver, thread_id="concurrent_writers")
+
+    n = 32
+    barrier = threading.Barrier(n)
+
+    def worker(i: int) -> None:
+        # Synchronize at the barrier to maximize contention.
+        barrier.wait()
+        saver.put_writes(
+            config=cp_config,
+            writes=[(f"channel_{i}", f"value_{i}")],
+            task_id=f"task_{i}",
+        )
+
+    with ThreadPoolExecutor(max_workers=n) as ex:
+        list(ex.map(worker, range(n)))
+
+    tpl = saver.get_tuple(cp_config)
+    assert tpl is not None
+    assert len(tpl.pending_writes) == n
+    seen_task_ids = {task_id for (task_id, _channel, _value) in tpl.pending_writes}
+    assert seen_task_ids == {f"task_{i}" for i in range(n)}
+
+
+def test_put_writes_retry_overwrites_in_place(saver) -> None:
+    """A retry of the same ``(task_id, idx)`` must upsert, not append."""
+    cp_config = _seed_checkpoint(saver, thread_id="retry_writer")
+
+    task_id = "retrying_task"
+
+    saver.put_writes(
+        config=cp_config,
+        writes=[("ch", "v1")],
+        task_id=task_id,
+    )
+    saver.put_writes(
+        config=cp_config,
+        writes=[("ch", "v2")],
+        task_id=task_id,
+    )
+    saver.put_writes(
+        config=cp_config,
+        writes=[("ch", "v3")],
+        task_id=task_id,
+    )
+
+    tpl = saver.get_tuple(cp_config)
+    assert tpl is not None
+
+    matching = [w for w in tpl.pending_writes if w[0] == task_id]
+    assert len(matching) == 1
+    assert matching[0][1] == "ch"
+    assert matching[0][2] == "v3"
+
+
+def test_put_concurrent_no_lost_timeline_entries(saver) -> None:
+    """N parallel ``put()`` calls must each leave a timeline entry."""
+    base_config: RunnableConfig = {
+        "configurable": {
+            "thread_id": "concurrent_timeline",
+            "checkpoint_ns": "",
+        }
+    }
+
+    n = 16
+    barrier = threading.Barrier(n)
+
+    def worker(i: int) -> None:
+        cp = empty_checkpoint()
+        cp["id"] = f"cp_{i:04d}"
+        cp["ts"] = f"2026-04-29T17:00:{i:02d}+00:00"
+        metadata = {"source": "input", "step": i, "writes": {}, "parents": {}}
+        barrier.wait()
+        saver.put(base_config, cp, metadata, {})  # type: ignore[arg-type]
+
+    with ThreadPoolExecutor(max_workers=n) as ex:
+        list(ex.map(worker, range(n)))
+
+    timeline = list(saver.list(base_config))
+    assert len(timeline) == n
+    seen_ids = {tpl.config["configurable"]["checkpoint_id"] for tpl in timeline}
+    assert seen_ids == {f"cp_{i:04d}" for i in range(n)}