diff --git a/ir/store.py b/ir/store.py index d8cee0c..7223e8f 100644 --- a/ir/store.py +++ b/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 diff --git a/tests/test_store_packed.py b/tests/test_store_packed.py index b255cbb..8d46467 100644 --- a/tests/test_store_packed.py +++ b/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