Detect a torn packed-cache read instead of serving wrong rows

ir/store.py

@@ -29,6 +29,7 @@
 from __future__ import annotations
 
 import copy
+import hashlib
 import io
 import json
 import os
@@ -45,6 +46,32 @@
 _PACKED_FORMAT = 1
 
 
+def _packed_content_sig(ids_json: bytes, metas_json: bytes) -> str:
+    """Signature binding a packed matrix to the exact ids/metas written with it.
+
+    Hashes the bytes of ``ids.json`` and ``metas.json`` (length-prefixed, so the
+    boundary between them is unambiguous). ``sig.json`` is written last and
+    carries this digest, so a ``sig``/``matrix`` pair from one writer cannot
+    validate another writer's ids/metas: the four packed files are written
+    independently, and a same-length mixture of two concurrent rebuilds passes
+    every length check while row *i* no longer belongs to ``ids[i]``.
+
+    >>> _packed_content_sig(b'["a"]', b"[{}]") == _packed_content_sig(
+    ...     b'["a"]', b"[{}]"
+    ... )
+    True
+    >>> _packed_content_sig(b'["a"]', b"[{}]") == _packed_content_sig(
+    ...     b'["b"]', b"[{}]"
+    ... )
+    False
+    """
+    digest = hashlib.sha256()
+    digest.update(len(ids_json).to_bytes(8, "big"))
+    digest.update(ids_json)
+    digest.update(metas_json)
+    return digest.hexdigest()
+
+
 def _ndarray_store(rootdir) -> MutableMapping[str, np.ndarray]:
     """A ``dol`` file store whose values are float32 ``ndarray``s."""
     import dol
@@ -387,19 +414,38 @@ def _load_packed(self):
             if sig.get("format") != _PACKED_FORMAT:
                 return None
             mat = np.load(paths["matrix"], mmap_mode="r")
-            ids = json.loads(paths["ids"].read_text(encoding="utf-8"))
-            metas = json.loads(paths["metas"].read_text(encoding="utf-8"))
+            ids_json = paths["ids"].read_bytes()
+            metas_json = paths["metas"].read_bytes()
+            ids = json.loads(ids_json)
+            metas = json.loads(metas_json)
         except (OSError, ValueError):
             return None
         if len(ids) != mat.shape[0] or len(metas) != len(ids):
             return None
+        # A cache written before ``sig.json`` carried these fields has neither;
+        # keep accepting it on the length checks alone rather than invalidating
+        # every existing cache. When they are present they must agree, or the
+        # set mixes two writers' files and reads as a miss (rebuild) instead of
+        # serving rows under the wrong ids.
+        content_sig = sig.get("content_sig")
+        if content_sig is not None and content_sig != _packed_content_sig(
+            ids_json, metas_json
+        ):
+            return None
+        shape = sig.get("shape")
+        if shape is not None and list(mat.shape) != list(shape):
+            return None
         return (ids, mat, metas)
 
     def _save_packed(self, result: tuple[list[str], np.ndarray, list[dict]]) -> None:
         """Persist a freshly built matrix to the packed cache (best-effort).
 
         Skips empty corpora. Writes ``sig.json`` last so a crash mid-write
-        leaves the cache marked invalid (no sig) rather than torn.
+        leaves the cache marked invalid (no sig) rather than torn. The sig also
+        carries a :func:`_packed_content_sig` digest of the ids/metas bytes and
+        the matrix shape, which is what lets :meth:`_load_packed` notice the
+        case a write order cannot defend against: a *second* writer replacing
+        some of the four files while this set is on disk.
         """
         if self._packed_dir is None:
             return
@@ -409,11 +455,21 @@ def _save_packed(self, result: tuple[list[str], np.ndarray, list[dict]]) -> None
         try:
             self._packed_dir.mkdir(parents=True, exist_ok=True)
             paths = self._packed_paths()
-            np.save(paths["matrix"], np.asarray(mat, dtype=np.float32))
-            paths["ids"].write_text(json.dumps(ids), encoding="utf-8")
-            paths["metas"].write_text(json.dumps(metas), encoding="utf-8")
+            arr = np.asarray(mat, dtype=np.float32)
+            ids_json = json.dumps(ids).encode("utf-8")
+            metas_json = json.dumps(metas).encode("utf-8")
+            np.save(paths["matrix"], arr)
+            paths["ids"].write_bytes(ids_json)
+            paths["metas"].write_bytes(metas_json)
             paths["sig"].write_text(
-                json.dumps({"format": _PACKED_FORMAT, "count": len(ids)}),
+                json.dumps(
+                    {
+                        "format": _PACKED_FORMAT,
+                        "count": len(ids),
+                        "content_sig": _packed_content_sig(ids_json, metas_json),
+                        "shape": list(arr.shape),
+                    }
+                ),
                 encoding="utf-8",
             )
             self._packed_stale = False

tests/test_store_packed.py

@@ -6,6 +6,8 @@
 and serving ``metas()`` without touching vectors (for lexical-only ranking).
 """
 
+import json
+
 import numpy as np
 
 from ir.base import Record
@@ -98,3 +100,70 @@ def test_memory_store_keeps_purely_in_process(tmp_path):
     assert store.matrix()[1].shape == (1, 3)
     ids, metas = store.metas()
     assert ids == ["r1"] and len(metas) == 1
+
+
+def test_torn_packed_cache_is_rejected_not_served(tmp_path):
+    """A matrix from one writer + ids/metas from another must never be served.
+
+    The four packed files are written independently, so two writers rebuilding a
+    same-size corpus can leave ``matrix.npy`` from A beside ``ids.json`` /
+    ``metas.json`` from B. Every length check still passes, but row *i* no
+    longer belongs to ``ids[i]`` — searches would answer confidently wrong.
+    """
+    store, root = _file_store(tmp_path)
+    store.put_record(_rec("r1", vec=(3.0, 0.0, 0.0)))
+    store.put_record(_rec("r2", vec=(0.0, 4.0, 0.0)))
+    store.matrix()  # builds from records + persists packed
+
+    packed = root / "matrix"
+    # Writer B rewrites ids/metas (same length, different order) after writer A
+    # wrote matrix.npy — the common same-size incremental case.
+    ids = json.loads((packed / "ids.json").read_text(encoding="utf-8"))
+    metas = json.loads((packed / "metas.json").read_text(encoding="utf-8"))
+    (packed / "ids.json").write_text(json.dumps(ids[::-1]), encoding="utf-8")
+    (packed / "metas.json").write_text(json.dumps(metas[::-1]), encoding="utf-8")
+
+    store2, _ = _file_store(tmp_path)
+    ids2, mat2, _metas2 = store2.matrix()
+    expected = {"r1": (1.0, 0.0, 0.0), "r2": (0.0, 1.0, 0.0)}
+    assert sorted(ids2) == ["r1", "r2"]
+    for i, rid in enumerate(ids2):
+        np.testing.assert_allclose(np.asarray(mat2[i]), expected[rid], atol=1e-6)
+
+
+def test_packed_cache_with_mismatched_matrix_shape_is_rejected(tmp_path):
+    """A matrix whose columns disagree with the recorded shape is a torn set."""
+    store, root = _file_store(tmp_path)
+    store.put_record(_rec("r1", vec=(3.0, 0.0, 0.0)))
+    store.put_record(_rec("r2", vec=(0.0, 4.0, 0.0)))
+    store.matrix()
+
+    # Writer B's matrix (a different embedding dim) over writer A's sig/ids.
+    np.save(root / "matrix" / "matrix.npy", np.zeros((2, 5), dtype=np.float32))
+
+    store2, _ = _file_store(tmp_path)
+    _ids2, mat2, _metas2 = store2.matrix()
+    assert mat2.shape == (2, 3)  # rebuilt from records, not the torn cache
+
+
+def test_packed_cache_without_content_sig_still_loads(tmp_path):
+    """A cache written by an older ``ir`` (no signature) stays usable as-is."""
+    store, root = _file_store(tmp_path)
+    store.put_record(_rec("r1", vec=(3.0, 0.0, 0.0)))
+    ids, _mat, metas = store.matrix()
+
+    sig_path = root / "matrix" / "sig.json"
+    sig = json.loads(sig_path.read_text(encoding="utf-8"))
+    sig.pop("content_sig", None)
+    sig.pop("shape", None)
+    sig_path.write_text(json.dumps(sig), encoding="utf-8")
+
+    store2, _ = _file_store(tmp_path)
+
+    def _no_rebuild():
+        raise AssertionError("legacy packed cache should be used, not rebuilt")
+
+    store2._build_matrix = _no_rebuild
+    ids2, _mat2, metas2 = store2.matrix()
+    assert ids2 == ids
+    assert metas2 == metas