From 16e744292dee7796b7272d0996fd6bd4fa1f9ac3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Fri, 14 Aug 2026 15:36:53 +0200 Subject: [PATCH] perf(retain): store-owned backend writes facts once with entities inline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A memories store that owns its rows (external backend) retains a document in a connection-free store phase. That phase wrote each memory TWICE: insert_facts staged it without entities, then record_unit_entities read the just-written records back (full vectors) and re-upserted them with entity ids attached — because entity ids can only be resolved onto real unit ids after they exist. On an object-store-backed engine that reattach is a second full write per memory plus a read-back, doubling the round-trips on the slow path (the dominant cost of retain). It's avoidable: mint the ids without writing (insert_facts_batch with defer_index), remap the entities, then write once via index_facts with the entity ids already inline — tagged with the write-group txn so it commits atomically with the group (and, as a bonus, the postings are now witness-covered instead of riding an uncovered seam). Co-occurrence accumulation still runs (record_unit_entity_postings gains store_write=False: co-occurrence only, since the row is already correct). Streaming ext path only; the delta path is unchanged. Tests updated: the single write via index_facts is asserted connection-free and txn-tagged, and the posting runs store_write=False. No behavior change for the Postgres store (a no-op there). --- .../hindsight_api/engine/entity_resolver.py | 37 ++++++++++++++----- .../engine/retain/fact_storage.py | 7 +++- .../engine/retain/orchestrator.py | 34 ++++++++++++++--- .../tests/test_retain_ext_writegroup.py | 35 ++++++++++++------ 4 files changed, 85 insertions(+), 28 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/entity_resolver.py b/hindsight-api-slim/hindsight_api/engine/entity_resolver.py index 3f80fb417..e5bab619b 100644 --- a/hindsight-api-slim/hindsight_api/engine/entity_resolver.py +++ b/hindsight-api-slim/hindsight_api/engine/entity_resolver.py @@ -1246,6 +1246,7 @@ async def record_unit_entity_postings( self, unit_entity_pairs: list[tuple[str, str]] | list[tuple[str, str, datetime | None]], bank_id: str | None = None, + store_write: bool = True, ): """Store-owned variant of :meth:`link_units_to_entities_batch` that touches NO Postgres connection. @@ -1257,6 +1258,13 @@ async def record_unit_entity_postings( its connection-free store phase and never hold the data-plane connection across the object-store write. NOT for the Postgres store, whose posting is a real ``unit_entities`` INSERT that requires the connection. + + ``store_write=False`` skips the store-side posting and does ONLY the co-occurrence + accumulation. The caller uses this when it has already attached entity ids to the memories + as part of the same write (a single deferred write with entities inline, instead of + write-then-reattach) — so the store row is already correct and a second store write would + be redundant. Co-occurrence still runs: it references only ``entities`` and is needed by the + entity-graph endpoint and resolution's disambiguation signal regardless of who wrote the row. """ if not unit_entity_pairs: return @@ -1265,10 +1273,16 @@ async def record_unit_entity_postings( (t[0], t[1], t[2] if len(t) >= 3 else None) # type: ignore[misc] for t in unit_entity_pairs ] - return await self._link_units_to_entities_batch_impl(None, normalized, bank_id) + return await self._link_units_to_entities_batch_impl( + None, normalized, bank_id, store_write=store_write + ) async def _link_units_to_entities_batch_impl( - self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]], bank_id: str | None = None + self, + conn, + unit_entity_pairs: list[tuple[str, str, datetime | None]], + bank_id: str | None = None, + store_write: bool = True, ): # Sorted bulk insert to prevent deadlocks from inconsistent lock ordering # across concurrent transactions on the unit_entities unique index. @@ -1280,16 +1294,19 @@ async def _link_units_to_entities_batch_impl( # memories store records it. Co-occurrence below is separate and unaffected: # it references only `entities`, which stays in Postgres either way, and is # read by the entity-graph endpoint and by resolution's disambiguation signal. + # `store_write=False` means the caller already wrote the postings inline with the + # memories, so we skip the (redundant) second store write and keep only co-occurrence. from .memories import get_memories - await get_memories().record_unit_entities( - conn=conn, - ops=self._ops, - fq_table=fq_table, - bank_id=bank_id, - unit_ids=unit_ids, - entity_ids=entity_ids, - ) + if store_write: + await get_memories().record_unit_entities( + conn=conn, + ops=self._ops, + fq_table=fq_table, + bank_id=bank_id, + unit_ids=unit_ids, + entity_ids=entity_ids, + ) # Build maps keyed by unit_id: # unit_to_entities: entity set per unit (for the co-occurrence cross-product) diff --git a/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py b/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py index d3523ad55..2fb4d115b 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py @@ -109,16 +109,21 @@ async def index_facts( facts: list[ProcessedFact], document_id: str | None = None, unit_entity_ids: dict[str, list[str]] | None = None, + txn=None, ) -> None: """Complete a deferred `insert_facts_batch`, now that the edges are known. ``unit_entity_ids`` is the unit→entity posting and each fact's causal relations are its edges; both travel with the memory for a store that owns them. A no-op for the Postgres store, which wrote all of it already. + + ``txn`` rides a cross-store write-group handle so this single, entity-bearing + write commits (and becomes visible) atomically with the rest of the group — + the store-owned retain path writes facts ONCE here rather than write-then-reattach. """ from ..memories import get_memories - await get_memories().index_facts(bank_id, unit_ids, facts, document_id, unit_entity_ids) + await get_memories().index_facts(bank_id, unit_ids, facts, document_id, unit_entity_ids, txn=txn) async def ensure_bank_exists(conn, bank_id: str, ops=None) -> None: diff --git a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py index a29487aff..29dc82809 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py @@ -670,14 +670,18 @@ async def _streaming_batch_write_ext( if cid: processed_fact.chunk_id = cid - # Stage the memory records to the store (conn unused by a store-owned backend), tagged with - # ext_txn. This is the slow object-store write we are keeping OUT of the connection window. - unit_ids = await fact_storage.insert_facts_batch(None, bank_id, batch_processed, ops=pool.ops, txn=ext_txn) + # Mint the unit ids WITHOUT writing (defer_index): entities can only be resolved onto real ids + # after they exist, so we write the memories to the store ONCE below — with their entity ids + # already attached — instead of writing them here and then re-upserting each one with entities. + # That reattach cost a SECOND full object-store write per memory (plus a read-back of the just- + # written records), doubling the store round-trips on the slow path. Connection-free either way. + unit_ids = await fact_storage.insert_facts_batch( + None, bank_id, batch_processed, ops=pool.ops, txn=ext_txn, defer_index=True + ) batch_result_ids = _map_results_to_contents(batch_contents, batch_processed, unit_ids if unit_ids else []) if unit_ids: - # Remap Phase-1 placeholder ids onto the real unit ids, then re-write each memory with its - # entity ids attached — also connection-free for a store-owned backend. + # Remap Phase-1 placeholder ids onto the real unit ids. resolved_entity_ids = [entity.entity_id for entity in phase1.entities.resolved_entities] remapped_entity_to_unit, _remapped_unit_to_entity_ids, _remapped_semantic = _remap_phase1_results( resolved_entity_ids, phase1.entities.entity_to_unit, phase1.entities.unit_to_entity_ids, [], unit_ids @@ -686,7 +690,25 @@ async def _streaming_batch_write_ext( (unit_id, resolved_entity_ids[idx], fact_date) for idx, (unit_id, _local_idx, fact_date) in enumerate(remapped_entity_to_unit) ] - await entity_resolver.record_unit_entity_postings(unit_entity_pairs, bank_id=bank_id) + # The single, entity-bearing store write — connection-free, and tagged with ext_txn so it + # commits (and becomes visible) atomically with the rest of the write-group. This replaces + # the earlier insert-then-reattach pair with one write. + unit_entity_ids: dict[str, list[str]] = {} + for unit_id, entity_id, _fd in unit_entity_pairs: + unit_entity_ids.setdefault(unit_id, []).append(entity_id) + await fact_storage.index_facts( + bank_id, + unit_ids, + batch_processed, + document_id=effective_doc_id, + unit_entity_ids=unit_entity_ids, + txn=ext_txn, + ) + # The store row was just written with its entities inline, so skip the (now redundant) + # second store write and keep ONLY the co-occurrence accumulation the entity graph needs. + await entity_resolver.record_unit_entity_postings( + unit_entity_pairs, bank_id=bank_id, store_write=False + ) # ---- CONNECTION PHASE (short transaction: local metadata + commit witness) ---- try: diff --git a/hindsight-api-slim/tests/test_retain_ext_writegroup.py b/hindsight-api-slim/tests/test_retain_ext_writegroup.py index e5486559b..e5a087b2b 100644 --- a/hindsight-api-slim/tests/test_retain_ext_writegroup.py +++ b/hindsight-api-slim/tests/test_retain_ext_writegroup.py @@ -63,12 +63,21 @@ def _make_common(monkeypatch, tracker, *, calls): """Patch the module-level collaborators the helper reaches for.""" monkeypatch.setattr(orch, "acquire_with_retry", lambda pool: tracker.acquire()) - async def _insert_facts_batch(conn, bank_id, processed, ops=None, txn=None): - calls.append(("insert_facts", conn, tracker.open)) - tracker.store_writes_saw_open.append(tracker.open) + async def _insert_facts_batch(conn, bank_id, processed, ops=None, txn=None, defer_index=False): + # defer_index=True mints ids WITHOUT writing — the entity-bearing write happens in + # index_facts below, so this call is a no-op store-side (records no write). + calls.append(("insert_facts", conn, tracker.open, defer_index)) assert conn is None, "ext store write must not receive a connection" + if not defer_index: + tracker.store_writes_saw_open.append(tracker.open) return ["u1"] + async def _index_facts(bank_id, unit_ids, facts, document_id=None, unit_entity_ids=None, txn=None): + # The single, entity-bearing store write — must be connection-free and carry the txn. + calls.append(("index_facts", tracker.open)) + tracker.store_writes_saw_open.append(tracker.open) + assert txn is not None, "the deferred store write must ride the write-group txn" + async def _store_chunks_batch(conn, bank_id, doc_id, meta, ops=None, store_document_text=True): calls.append(("store_chunks", tracker.open)) return {} @@ -77,6 +86,7 @@ async def _handle_doc_tracking(conn, *a, **k): calls.append(("handle_doc_tracking", tracker.open)) monkeypatch.setattr(orch.fact_storage, "insert_facts_batch", _insert_facts_batch) + monkeypatch.setattr(orch.fact_storage, "index_facts", _index_facts) monkeypatch.setattr(orch.chunk_storage, "store_chunks_batch", _store_chunks_batch) monkeypatch.setattr(orch.fact_storage, "handle_document_tracking", _handle_doc_tracking) monkeypatch.setattr(orch, "_map_results_to_contents", lambda contents, pf, uids: [list(uids)]) @@ -109,10 +119,11 @@ def __init__(self, tracker): self.postings = [] self.reasserts = [] - async def record_unit_entity_postings(self, pairs, bank_id=None): - # THE contract: the store re-posting runs with no connection held. + async def record_unit_entity_postings(self, pairs, bank_id=None, store_write=True): + # THE contract: co-occurrence accumulation runs with no connection held. store_write=False + # here because the entity ids were already written inline by index_facts (single write). assert self._t.open is False, "entity posting must run connection-free" - self.postings.append(pairs) + self.postings.append((pairs, store_write)) async def reassert_entities_batch(self, bank_id, resolved, conn): assert conn is not None @@ -174,10 +185,11 @@ async def test_store_writes_are_connection_free_and_witness_is_in_txn(monkeypatc assert result.aborted is False assert result.batch_result_ids == [["u1"]] - # The fact write happened with no connection held. + # The fact write happened with no connection held (a single deferred write via index_facts). assert tracker.store_writes_saw_open == [False] - # Entity re-posting happened (also connection-free — asserted inside the fake). - assert er.postings == [[("u1", "e1", None)]] + assert ("index_facts", False) in calls # the entity-bearing store write ran connection-free + # Co-occurrence ran connection-free, with store_write=False (entities already written inline). + assert er.postings == [([("u1", "e1", None)], False)] # Witness written with a real connection, exactly once; commit published after release. assert len(provider.witnesses) == 1 and provider.witnesses[0][1] is not None assert provider.decisions == [True] @@ -333,10 +345,11 @@ async def test_delta_store_writes_are_connection_free(monkeypatch): assert result.fell_back is False assert result.result_unit_ids == [["u1"]] - # Fact write + body store + entity re-posting all happened connection-free. + # Fact write + body store + entity re-posting all happened connection-free. The delta path + # still writes-then-reposts (store_write=True) — only the streaming path uses the single write. assert tracker.store_writes_saw_open == [False] assert ("store_document_bodies", False) in calls - assert er.postings == [[("u1", "e1", None)]] + assert er.postings == [([("u1", "e1", None)], True)] # Witness inside the txn; publish after release. assert len(provider.witnesses) == 1 and provider.witnesses[0][1] is not None assert provider.decisions == [True]