holetron-lab/fleet-memory
1
0
Fork 0
Code Releases Packages Activity Actions

fix(retain): run _ann_seeds temp table inside a transaction (#954)

Browse source 
compute_semantic_links_ann created a TEMP TABLE outside any transaction,
then ran a TRUNCATE / COPY / SELECT / DROP sequence as separate statements
on the same asyncpg connection. This is fine against a direct Postgres
connection but fails intermittently when the caller is routed through
PgBouncer in transaction pool mode:

  CREATE TEMP TABLE IF NOT EXISTS _ann_seeds (...)   -- backend A
  TRUNCATE _ann_seeds                                 -- backend B -> FAILS

Temp tables are session-scoped to the backend that created them. In
PgBouncer transaction mode the backend is only pinned to the client for
the duration of an actual transaction, so between standalone statements
the pooler can (and under concurrency, will) rebind the client to a
different backend. When that happens the _ann_seeds table disappears
and the follow-up statement fails with:

  relation "_ann_seeds" does not exist

Symptom: ~3% of sync retain calls (2 of 61) failed the Hindsight Cloud
smoke test on a recent hindsight-dev deploy. Async retains are masked
by the 3-attempt retry loop so they usually eventually succeed.

Fix: wrap the CREATE TEMP TABLE -> COPY -> SELECT sequence in a single
`async with conn.transaction():` block, and use ON COMMIT DROP so the
temp table is transaction-scoped and auto-cleaned at commit. Also
switch `SET hnsw.ef_search = 60` to `SET LOCAL` so the tuning is
transaction-scoped and no longer leaks onto the pooled backend for
subsequent recall queries. Drop the now-unnecessary manual TRUNCATE,
explicit DROP TABLE, and RESET hnsw.ef_search.

The function docstring still correctly describes this as running on a
separate connection outside the surrounding write transaction — this
change only adds an inner transaction around the ANN work itself to
keep the temp table visible to PgBouncer.

Tests:
- Add TestComputeSemanticLinksAnnPgBouncerSafety with 5 regression
  tests using a mocked connection. These are structural asserts — they
  check that the function enters conn.transaction(), uses ON COMMIT DROP,
  uses SET LOCAL, and does not reintroduce manual TRUNCATE / DROP /
  RESET calls. They would have caught the original bug if they had
  existed, and will catch any future reversion.
This commit is contained in:
Chris Bartholomew 2026-04-10 12:36:03 -04:00 committed by GitHub
parent e22ae05f47
commit 3fc87e767c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 209 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

50
hindsight-api-slim/hindsight_api/engine/retain/link_utils.py
View file

@ -739,17 +739,10 @@ async def compute_semantic_links_ann(
return []
import time as time_mod
import uuid as uuid_mod
ann_start = time_mod.time()
links = []
# Lower ef_search for retain ANN — default 400 is tuned for recall precision
# but at 164k units each HNSW probe takes 94ms. ef_search=60 gives 2.7ms/probe
# (35x faster) with sufficient accuracy for top-50 semantic link creation.
# Reset after to avoid polluting the connection pool for recall queries.
await conn.execute("SET hnsw.ef_search = 60")
logger.debug(f"[ANN] Starting: {len(unit_ids)} seeds, top_k={top_k}")
# Build per-unit fact_types (default to 'world' if not provided)
@ -760,18 +753,43 @@ async def compute_semantic_links_ann(
# sequential-scan every HNSW probe result against the array, destroying
# performance (67s for 8k seeds). Self-links are harmless (ON CONFLICT DO
# NOTHING handles duplicates in memory_links).
#
# The entire CREATE TEMP TABLE → COPY → SELECT sequence MUST run inside a
# single transaction. Callers may connect through pgBouncer in `transaction`
# pool mode, in which case the backend is only pinned to the client for the
# duration of a transaction. Outside a transaction, pgBouncer can rebind
# the client to a different backend between statements, and the temp table
# (which is session-scoped to its creating backend) becomes invisible.
# The observed failure mode was an intermittent
# `relation "_ann_seeds" does not exist` on the second statement.
#
# Using ON COMMIT DROP + SET LOCAL also means we don't have to remember to
# manually drop the temp table or reset hnsw.ef_search — the transaction
# end handles both.
rows: list = []
async with conn.transaction():
# Transaction-local ef_search. Default 400 is tuned for recall precision
# but at 164k units each HNSW probe takes 94ms. ef_search=60 gives 2.7ms
# per probe (35x faster) with sufficient accuracy for top-50 semantic
# link creation. SET LOCAL auto-reverts at commit, so we don't pollute
# the pool for subsequent recall queries.
await conn.execute("SET LOCAL hnsw.ef_search = 60")
t_setup = time_mod.time()
await conn.execute("CREATE TEMP TABLE IF NOT EXISTS _ann_seeds (unit_id text, emb_text text, fact_type text)")
await conn.execute("TRUNCATE _ann_seeds")
await conn.execute(
"CREATE TEMP TABLE _ann_seeds (unit_id text, emb_text text, fact_type text) ON COMMIT DROP"
)
records = [
(uid, emb if isinstance(emb, str) else str(emb), ft) for uid, emb, ft in zip(unit_ids, embeddings, fact_types)
(uid, emb if isinstance(emb, str) else str(emb), ft)
for uid, emb, ft in zip(unit_ids, embeddings, fact_types)
]
await conn.copy_records_to_table("_ann_seeds", records=records, columns=["unit_id", "emb_text", "fact_type"])
await conn.copy_records_to_table(
"_ann_seeds", records=records, columns=["unit_id", "emb_text", "fact_type"]
)
logger.debug(f"[ANN] Temp table setup: {time_mod.time() - t_setup:.3f}s ({len(records)} seeds)")
# Run one ANN query per fact_type so each uses the right HNSW index.
rows = []
active_types = set(fact_types)
for fact_type in active_types:
t_query = time_mod.time()
@ -802,12 +820,8 @@ async def compute_semantic_links_ann(
)
logger.debug(f"[ANN] fact_type={fact_type}: {len(ft_rows)} rows in {time_mod.time() - t_query:.3f}s")
rows.extend(ft_rows)
# Clean up temp table (no ON COMMIT DROP since we're not in a transaction)
await conn.execute("DROP TABLE IF EXISTS _ann_seeds")
# Reset ef_search to default so the pooled connection doesn't affect recall queries
await conn.execute("RESET hnsw.ef_search")
# Transaction commits here. _ann_seeds is dropped (ON COMMIT DROP).
# hnsw.ef_search reverts (SET LOCAL).
for row in rows:
sim = float(min(1.0, max(0.0, row["similarity"])))

144
hindsight-api-slim/tests/test_link_utils.py
View file

@ -2,12 +2,14 @@
import numpy as np
import pytest
from datetime import datetime, timezone, timedelta
from unittest.mock import AsyncMock, MagicMock
from hindsight_api.engine.retain.link_utils import (
_normalize_datetime,
_cap_links_per_unit,
compute_temporal_links,
compute_temporal_query_bounds,
compute_semantic_links_ann,
compute_semantic_links_within_batch,
MAX_TEMPORAL_LINKS_PER_UNIT,
)
@ -388,3 +390,145 @@ class TestComputeSemanticLinksWithinBatch:
assert link_type == "semantic"
assert 0.0 <= weight <= 1.0
assert entity_id is None
class TestComputeSemanticLinksAnnPgBouncerSafety:
"""Regression tests ensuring compute_semantic_links_ann stays in a single
transaction so that the `_ann_seeds` temp table remains visible when the
caller's connection goes through pgBouncer in `transaction` pool mode.
In pgBouncer transaction mode, the backend is only pinned to the client
for the duration of an actual PostgreSQL transaction. Outside a
transaction, consecutive statements can land on different backends, and
session-scoped temp tables (which are bound to the backend that created
them) become invisible. The observed failure mode was an intermittent
`relation "_ann_seeds" does not exist` on the statement immediately
following the CREATE TEMP TABLE.
"""
@pytest.fixture
def mock_conn(self):
"""An asyncpg-like connection mock with an async `transaction()`
context manager and awaitable execute/fetch/copy helpers."""
conn = MagicMock()
txn_cm = MagicMock()
txn_cm.__aenter__ = AsyncMock(return_value=None)
txn_cm.__aexit__ = AsyncMock(return_value=None)
conn.transaction = MagicMock(return_value=txn_cm)
conn.execute = AsyncMock()
conn.copy_records_to_table = AsyncMock()
conn.fetch = AsyncMock(return_value=[])
return conn
@pytest.mark.asyncio
async def test_empty_inputs_skip_transaction(self, mock_conn):
"""No seeds -> no work, no transaction, no temp-table churn."""
result = await compute_semantic_links_ann(
conn=mock_conn,
bank_id="bank-1",
unit_ids=[],
embeddings=[],
)
assert result == []
mock_conn.transaction.assert_not_called()
mock_conn.execute.assert_not_called()
@pytest.mark.asyncio
async def test_runs_inside_a_transaction(self, mock_conn):
"""The full CREATE TEMP TABLE -> COPY -> SELECT sequence must happen
inside a single `async with conn.transaction():` block."""
emb = [0.1] * 384
await compute_semantic_links_ann(
conn=mock_conn,
bank_id="bank-1",
unit_ids=["u1", "u2"],
embeddings=[emb, emb],
fact_types=["world", "world"],
)
# Transaction context manager was entered.
mock_conn.transaction.assert_called_once()
txn_cm = mock_conn.transaction.return_value
txn_cm.__aenter__.assert_awaited_once()
txn_cm.__aexit__.assert_awaited_once()
@pytest.mark.asyncio
async def test_temp_table_uses_on_commit_drop(self, mock_conn):
"""The CREATE TEMP TABLE statement must use ON COMMIT DROP so the
table is transaction-scoped. Without ON COMMIT DROP the table would
be session-scoped and would not survive pgBouncer backend rebinding
between transactions."""
emb = [0.1] * 384
await compute_semantic_links_ann(
conn=mock_conn,
bank_id="bank-1",
unit_ids=["u1"],
embeddings=[emb],
fact_types=["world"],
)
executed_sql = [call.args[0] for call in mock_conn.execute.call_args_list]
create_statements = [s for s in executed_sql if "CREATE TEMP TABLE" in s]
assert len(create_statements) == 1, "Should create _ann_seeds exactly once"
assert "_ann_seeds" in create_statements[0]
assert "ON COMMIT DROP" in create_statements[0], (
"CREATE TEMP TABLE must use ON COMMIT DROP so the table is cleaned "
"up at transaction end and is transaction-scoped"
)
# Must not use IF NOT EXISTS — the table is fresh each transaction.
assert "IF NOT EXISTS" not in create_statements[0], (
"With ON COMMIT DROP the table is always fresh at transaction start, "
"so IF NOT EXISTS is both unnecessary and misleading (suggests the "
"table might persist across transactions)"
)
@pytest.mark.asyncio
async def test_no_manual_drop_or_truncate(self, mock_conn):
"""With ON COMMIT DROP we must not re-add manual TRUNCATE or DROP
statements — they were the source of the original pgBouncer bug."""
emb = [0.1] * 384
await compute_semantic_links_ann(
conn=mock_conn,
bank_id="bank-1",
unit_ids=["u1"],
embeddings=[emb],
fact_types=["world"],
)
executed_sql = [call.args[0] for call in mock_conn.execute.call_args_list]
assert not any("TRUNCATE _ann_seeds" in s for s in executed_sql), (
"TRUNCATE is unnecessary with ON COMMIT DROP and was previously "
"the statement that failed with 'relation does not exist' when "
"pgBouncer rebound the backend"
)
assert not any("DROP TABLE" in s and "_ann_seeds" in s for s in executed_sql), (
"Explicit DROP is unnecessary with ON COMMIT DROP"
)
@pytest.mark.asyncio
async def test_uses_set_local_for_ef_search(self, mock_conn):
"""hnsw.ef_search must be set with SET LOCAL so the change is scoped
to the transaction. Without SET LOCAL, the setting would leak onto
the pooled backend and affect subsequent recall queries that land
on the same backend."""
emb = [0.1] * 384
await compute_semantic_links_ann(
conn=mock_conn,
bank_id="bank-1",
unit_ids=["u1"],
embeddings=[emb],
fact_types=["world"],
)
executed_sql = [call.args[0] for call in mock_conn.execute.call_args_list]
ef_statements = [s for s in executed_sql if "hnsw.ef_search" in s]
assert ef_statements, "ef_search must be tuned down for retain ANN"
for stmt in ef_statements:
assert stmt.strip().startswith("SET LOCAL"), (
f"hnsw.ef_search must use SET LOCAL, got: {stmt}"
)
# And there must not be a RESET — SET LOCAL handles it at commit.
assert not any("RESET hnsw.ef_search" in s for s in executed_sql)