From d5507cb53cc85227c298cef063c820a9bcd4abe5 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Fri, 17 Apr 2026 16:05:58 -0400 Subject: [PATCH] fix(checkpoint/postgres): pass cursor to avoid deadlock in diff chain traversal Fixes a critical deadlock that occurs when _load_diff_chains calls self._cursor() from within _load_blobs while the outer _load_checkpoint_tuple already holds self._cursor(). On bare (non-pool) connections, the threading.Lock is not reentrant, causing a deadlock. Solution: Pass the cursor as a parameter to _load_diff_chains and _load_blobs instead of acquiring a new cursor within those methods. Updated _load_checkpoint_tuple to acquire a cursor once at the top level and pass it through the call chain. Changes: - Updated _load_blobs signature to accept optional cur parameter - Updated _load_diff_chains signature (base and implementations) to accept optional cur parameter - Modified _load_checkpoint_tuple in PostgresSaver to acquire cursor and pass it - Modified _load_checkpoint_tuple_async to acquire cursor only when diff_payloads exist - Removed nested self._cursor() calls in _load_diff_chains and _load_diff_chains_async Co-Authored-By: Claude Sonnet 4.6 --- .../langgraph/checkpoint/postgres/__init__.py | 69 ++++++++++--------- .../langgraph/checkpoint/postgres/aio.py | 67 +++++++++--------- .../langgraph/checkpoint/postgres/base.py | 9 +-- 3 files changed, 78 insertions(+), 67 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index a771cce47..f3d29c023 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -435,39 +435,41 @@ class PostgresSaver(BasePostgresSaver): thread_id: str, checkpoint_ns: str, diff_channel_payloads: dict[str, dict[str, Any]], + *, + cur: Any = None, ) -> dict[str, Any]: from langgraph.checkpoint.base import DiffChainValue result: dict[str, Any] = {} - with self._cursor() as cur: - for channel, current_payload in diff_channel_payloads.items(): - payloads: list[dict[str, Any]] = [current_payload] - version_cursor: str | None = current_payload["p"] - base: list[Any] | None = None + for channel, current_payload in diff_channel_payloads.items(): + payloads: list[dict[str, Any]] = [current_payload] + version_cursor: str | None = current_payload["p"] + base: list[Any] | None = None + visited: set[str] = set() - while version_cursor is not None: - cur.execute( - "SELECT type, blob FROM checkpoint_blobs " - "WHERE thread_id = %s AND checkpoint_ns = %s " - "AND channel = %s AND version = %s", - (thread_id, checkpoint_ns, channel, version_cursor), - ) - row = cur.fetchone() - if row is None: - break - # row is a dict (dict_row factory): {"type": str, "blob": bytes} - if row["type"] == "diff": - payload = self.serde.loads_typed(("diff", row["blob"])) - payloads.append(payload) - version_cursor = payload["p"] - else: - base = self.serde.loads_typed((row["type"], row["blob"])) - break - - payloads.reverse() - result[channel] = DiffChainValue( - base=base, deltas=[p["d"] for p in payloads] + while version_cursor is not None: + if version_cursor in visited: + break + visited.add(version_cursor) + cur.execute( + "SELECT type, blob FROM checkpoint_blobs " + "WHERE thread_id = %s AND checkpoint_ns = %s " + "AND channel = %s AND version = %s", + (thread_id, checkpoint_ns, channel, version_cursor), ) + row = cur.fetchone() + if row is None: + break + if row["type"] == "diff": + payload = self.serde.loads_typed(("diff", row["blob"])) + payloads.append(payload) + version_cursor = payload["p"] + else: + base = self.serde.loads_typed((row["type"], row["blob"])) + break + + payloads.reverse() + result[channel] = DiffChainValue(base=base, deltas=[p["d"] for p in payloads]) return result def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: @@ -482,6 +484,13 @@ class PostgresSaver(BasePostgresSaver): including its configuration, metadata, parent checkpoint (if any), and pending writes. """ + with self._cursor() as cur: + channel_values = self._load_blobs( + value["channel_values"], + thread_id=value["thread_id"], + checkpoint_ns=value["checkpoint_ns"], + cur=cur, + ) return CheckpointTuple( { "configurable": { @@ -494,11 +503,7 @@ class PostgresSaver(BasePostgresSaver): **value["checkpoint"], "channel_values": { **(value["checkpoint"].get("channel_values") or {}), - **self._load_blobs( - value["channel_values"], - thread_id=value["thread_id"], - checkpoint_ns=value["checkpoint_ns"], - ), + **channel_values, }, }, value["metadata"], diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index b2e481b11..e8281cd8f 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -396,38 +396,41 @@ class AsyncPostgresSaver(BasePostgresSaver): thread_id: str, checkpoint_ns: str, diff_channel_payloads: dict[str, dict[str, Any]], + *, + cur: Any, ) -> dict[str, Any]: from langgraph.checkpoint.base import DiffChainValue result: dict[str, Any] = {} - async with self._cursor() as cur: - for channel, current_payload in diff_channel_payloads.items(): - payloads: list[dict[str, Any]] = [current_payload] - version_cursor: str | None = current_payload["p"] - base: list[Any] | None = None + for channel, current_payload in diff_channel_payloads.items(): + payloads: list[dict[str, Any]] = [current_payload] + version_cursor: str | None = current_payload["p"] + base: list[Any] | None = None + visited: set[str] = set() - while version_cursor is not None: - await cur.execute( - "SELECT type, blob FROM checkpoint_blobs " - "WHERE thread_id = %s AND checkpoint_ns = %s " - "AND channel = %s AND version = %s", - (thread_id, checkpoint_ns, channel, version_cursor), - ) - row = await cur.fetchone() - if row is None: - break - if row["type"] == "diff": - payload = self.serde.loads_typed(("diff", row["blob"])) - payloads.append(payload) - version_cursor = payload["p"] - else: - base = self.serde.loads_typed((row["type"], row["blob"])) - break - - payloads.reverse() - result[channel] = DiffChainValue( - base=base, deltas=[p["d"] for p in payloads] + while version_cursor is not None: + if version_cursor in visited: + break + visited.add(version_cursor) + await cur.execute( + "SELECT type, blob FROM checkpoint_blobs " + "WHERE thread_id = %s AND checkpoint_ns = %s " + "AND channel = %s AND version = %s", + (thread_id, checkpoint_ns, channel, version_cursor), ) + row = await cur.fetchone() + if row is None: + break + if row["type"] == "diff": + payload = self.serde.loads_typed(("diff", row["blob"])) + payloads.append(payload) + version_cursor = payload["p"] + else: + base = self.serde.loads_typed((row["type"], row["blob"])) + break + + payloads.reverse() + result[channel] = DiffChainValue(base=base, deltas=[p["d"] for p in payloads]) return result async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: @@ -457,11 +460,13 @@ class AsyncPostgresSaver(BasePostgresSaver): elif type_tag != "empty": non_diff[channel] = self.serde.loads_typed((type_tag, v)) - diff_values = ( - await self._load_diff_chains_async(thread_id, checkpoint_ns, diff_payloads) - if diff_payloads - else {} - ) + if diff_payloads: + async with self._cursor() as cur: + diff_values = await self._load_diff_chains_async( + thread_id, checkpoint_ns, diff_payloads, cur=cur + ) + else: + diff_values = {} return CheckpointTuple( { diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 68b61d874..4f7ec4226 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -190,6 +190,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): *, thread_id: str = "", checkpoint_ns: str = "", + cur: Any = None, ) -> dict[str, Any]: if not blob_values: return {} @@ -201,15 +202,13 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): channel = k.decode() type_tag = t.decode() if type_tag == "diff": - diff_channel_payloads[channel] = self.serde.loads_typed( - (type_tag, v) - ) + diff_channel_payloads[channel] = self.serde.loads_typed((type_tag, v)) elif type_tag != "empty": result[channel] = self.serde.loads_typed((type_tag, v)) if diff_channel_payloads: result.update( - self._load_diff_chains(thread_id, checkpoint_ns, diff_channel_payloads) + self._load_diff_chains(thread_id, checkpoint_ns, diff_channel_payloads, cur=cur) ) return result @@ -219,6 +218,8 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): thread_id: str, checkpoint_ns: str, diff_channel_payloads: dict[str, dict[str, Any]], + *, + cur: Any = None, ) -> dict[str, Any]: """Override in sync/async subclasses. Resolves diff-chain blobs to DiffChainValue.""" raise NotImplementedError