fix(retain): run _ann_seeds temp table inside a transaction (#954)
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:
parent
e22ae05f47
commit
3fc87e767c
2 changed files with 209 additions and 51 deletions
|
|
@ -739,17 +739,10 @@ async def compute_semantic_links_ann(
|
||||||
return []
|
return []
|
||||||
|
|
||||||
import time as time_mod
|
import time as time_mod
|
||||||
import uuid as uuid_mod
|
|
||||||
|
|
||||||
ann_start = time_mod.time()
|
ann_start = time_mod.time()
|
||||||
links = []
|
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}")
|
logger.debug(f"[ANN] Starting: {len(unit_ids)} seeds, top_k={top_k}")
|
||||||
|
|
||||||
# Build per-unit fact_types (default to 'world' if not provided)
|
# 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
|
# sequential-scan every HNSW probe result against the array, destroying
|
||||||
# performance (67s for 8k seeds). Self-links are harmless (ON CONFLICT DO
|
# performance (67s for 8k seeds). Self-links are harmless (ON CONFLICT DO
|
||||||
# NOTHING handles duplicates in memory_links).
|
# 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()
|
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(
|
||||||
await conn.execute("TRUNCATE _ann_seeds")
|
"CREATE TEMP TABLE _ann_seeds (unit_id text, emb_text text, fact_type text) ON COMMIT DROP"
|
||||||
|
)
|
||||||
|
|
||||||
records = [
|
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)")
|
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.
|
# Run one ANN query per fact_type so each uses the right HNSW index.
|
||||||
rows = []
|
|
||||||
active_types = set(fact_types)
|
active_types = set(fact_types)
|
||||||
for fact_type in active_types:
|
for fact_type in active_types:
|
||||||
t_query = time_mod.time()
|
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")
|
logger.debug(f"[ANN] fact_type={fact_type}: {len(ft_rows)} rows in {time_mod.time() - t_query:.3f}s")
|
||||||
rows.extend(ft_rows)
|
rows.extend(ft_rows)
|
||||||
|
# Transaction commits here. _ann_seeds is dropped (ON COMMIT DROP).
|
||||||
# Clean up temp table (no ON COMMIT DROP since we're not in a transaction)
|
# hnsw.ef_search reverts (SET LOCAL).
|
||||||
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")
|
|
||||||
|
|
||||||
for row in rows:
|
for row in rows:
|
||||||
sim = float(min(1.0, max(0.0, row["similarity"])))
|
sim = float(min(1.0, max(0.0, row["similarity"])))
|
||||||
|
|
|
||||||
|
|
@ -2,12 +2,14 @@
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import pytest
|
import pytest
|
||||||
from datetime import datetime, timezone, timedelta
|
from datetime import datetime, timezone, timedelta
|
||||||
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
|
||||||
from hindsight_api.engine.retain.link_utils import (
|
from hindsight_api.engine.retain.link_utils import (
|
||||||
_normalize_datetime,
|
_normalize_datetime,
|
||||||
_cap_links_per_unit,
|
_cap_links_per_unit,
|
||||||
compute_temporal_links,
|
compute_temporal_links,
|
||||||
compute_temporal_query_bounds,
|
compute_temporal_query_bounds,
|
||||||
|
compute_semantic_links_ann,
|
||||||
compute_semantic_links_within_batch,
|
compute_semantic_links_within_batch,
|
||||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||||
)
|
)
|
||||||
|
|
@ -388,3 +390,145 @@ class TestComputeSemanticLinksWithinBatch:
|
||||||
assert link_type == "semantic"
|
assert link_type == "semantic"
|
||||||
assert 0.0 <= weight <= 1.0
|
assert 0.0 <= weight <= 1.0
|
||||||
assert entity_id is None
|
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)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue