From de5a781197ac6eee420fe8e31d6ef798401fb1d8 Mon Sep 17 00:00:00 2001 From: Andrei Marinescu Date: Wed, 2 Sep 2026 10:52:20 +0300 Subject: [PATCH 1/3] =?UTF-8?q?hnsw:=20URL=E2=86=92node=20index,=20O(k)=20?= =?UTF-8?q?reclaim,=20dirty-set=20incremental=20persist,=20slot-swapped=20?= =?UTF-8?q?full=20persist,=20zombie-transit=20traversal?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-URL node index makes MarkURLPassagesInvalid O(k) and LookupVectorByURL O(1) (live nodes only); valid/zombie counts are maintained, not scanned. Incremental persists now rewrite every node dirtied since the last persist (invalidated, or handed a back-link), so reclaim and reconcile survive a restart and shutdown no longer needs a full persist. Persist encodes per window under the read lock instead of holding it for the whole write. Full persists write into the inactive node slot ('v'+0x02) and only then point meta at it (HSW2 carries the slot; HSW1 is still emitted for slot 0x01). The previous graph stays loadable throughout and the two key ranges never share tombstones. searchLayer treats zombies as transit-only (traversed, never admitted to the ef window) and invalidating the entry point relocates it to the highest-level live node. --- cmd/cosift/async_hnsw_load_test.go | 2 +- cmd/cosift/serve_setup.go | 10 +- internal/index/hnsw.go | 183 ++++++++----- internal/index/hnsw_compact.go | 12 +- internal/index/hnsw_compact_test.go | 14 +- internal/index/hnsw_index_test.go | 315 +++++++++++++++++++++++ internal/index/hnsw_persist.go | 224 +++++++++++----- internal/index/hnsw_persist_test.go | 4 +- internal/index/hnsw_swap_test.go | 327 ++++++++++++++++++++++++ internal/store/pebble.go | 81 ++++-- internal/store/pebble_uncovered_test.go | 48 +++- 11 files changed, 1063 insertions(+), 157 deletions(-) create mode 100644 internal/index/hnsw_index_test.go create mode 100644 internal/index/hnsw_swap_test.go diff --git a/cmd/cosift/async_hnsw_load_test.go b/cmd/cosift/async_hnsw_load_test.go index c45be25..a23e552 100644 --- a/cmd/cosift/async_hnsw_load_test.go +++ b/cmd/cosift/async_hnsw_load_test.go @@ -52,7 +52,7 @@ func TestLoadHNSWProgressSkipsCorruptNode(t *testing.T) { t.Fatalf("persist: %v", err) } // Overwrite one node's blob with garbage too short to decode. - if err := f.ps.PutVectorNode(context.Background(), 1, []byte{0x00, 0x01, 0x02}); err != nil { + if err := f.ps.PutVectorNode(context.Background(), store.VectorSlotA, 1, []byte{0x00, 0x01, 0x02}); err != nil { t.Fatalf("corrupt node: %v", err) } g, ok, err := index.LoadHNSW(context.Background(), f.ps) diff --git a/cmd/cosift/serve_setup.go b/cmd/cosift/serve_setup.go index 2ea17d5..ad01f3f 100644 --- a/cmd/cosift/serve_setup.go +++ b/cmd/cosift/serve_setup.go @@ -77,10 +77,12 @@ func runPebbleServe(ctx context.Context, cfg *config.Config, args []string) erro vectorDim = meta.Dim vectorNodes = meta.NodeCount } else { - _ = ps.IterateVectorNodes(ctx, func(_ uint64, _ []byte) bool { - hasVectors = true - return false - }) + for _, slot := range []byte{store.VectorSlotA, store.VectorSlotB} { + _ = ps.IterateVectorNodes(ctx, slot, func(_ uint64, _ []byte) bool { + hasVectors = true + return false + }) + } } // The full graph is loaded asynchronously (loadHNSWInto, launched after // the listener binds) so the O(N) decode of millions of nodes no longer diff --git a/internal/index/hnsw.go b/internal/index/hnsw.go index e7c6191..f477d90 100644 --- a/internal/index/hnsw.go +++ b/internal/index/hnsw.go @@ -25,6 +25,9 @@ import ( "math/rand" "sort" "sync" + "sync/atomic" + + "github.com/pilot-protocol/cosift/internal/store" ) // HNSW default parameters. Tunable per build; suitable for general-purpose @@ -71,6 +74,13 @@ type HNSW struct { // raw vectors; new nodes need a subsequent pq-train to get codes. codebook *PQCodebook codes [][]uint16 // parallel to nodes; nil entries fall back to raw vec + + byURL map[string][]int32 // live node ids per URL + valid int // nodes with a vector + reclaimed atomic.Uint64 + dirty map[int32]struct{} // nodes changed since the last persist (vec cleared or neighbors gained) + persistMu sync.Mutex // serializes persist and compact + slot byte // on-disk node slot the graph was loaded from / last persisted to } type hnswNode struct { @@ -91,9 +101,23 @@ func NewHNSW(dim int) *HNSW { efSearch: HNSWefSearch, levelMult: 1.0 / math.Log(float64(HNSWM)), entryPoint: -1, + byURL: make(map[string][]int32), + dirty: make(map[int32]struct{}), + slot: store.VectorSlotA, } } +// Reclaimed reports the cumulative count of nodes invalidated by +// MarkURLPassagesInvalid. +func (h *HNSW) Reclaimed() uint64 { return h.reclaimed.Load() } + +// Slot reports the on-disk node slot the graph is currently persisted in. +func (h *HNSW) Slot() byte { + h.mu.RLock() + defer h.mu.RUnlock() + return h.slot +} + // SetEfSearch overrides the query-time candidate-list size. Bigger values // raise recall at proportional cost. Exposed for runtime tuning // (env COSIFT_HNSW_EF_SEARCH) after we observed Recall@10 dropping to ~0.47 @@ -156,22 +180,18 @@ type PQStatus struct { func (h *HNSW) PQStatus() PQStatus { h.mu.RLock() defer h.mu.RUnlock() - st := PQStatus{NodesTotal: len(h.nodes)} - if h.codebook != nil { - st.Enabled = true - st.Dim = h.codebook.Dim - st.M = h.codebook.M - st.K = h.codebook.K - } - // Walk in a single pass; count valid vecs and codes that are - // ATTACHED to those valid vecs (ghost codes on vec-less zombies - // don't help anyone search and shouldn't pad coverage). - for i := range h.nodes { - valid := len(h.nodes[i].vec) > 0 - if valid { - st.NodesValid++ - } - if h.codebook != nil && i < len(h.codes) && len(h.codes[i]) == h.codebook.M && valid { + st := PQStatus{NodesTotal: len(h.nodes), NodesValid: h.valid} + if h.codebook == nil { + return st + } + st.Enabled = true + st.Dim = h.codebook.Dim + st.M = h.codebook.M + st.K = h.codebook.K + // Count only codes ATTACHED to valid vecs (ghost codes on vec-less + // zombies don't help anyone search and shouldn't pad coverage). + for i := range h.codes { + if i < len(h.nodes) && len(h.nodes[i].vec) > 0 && len(h.codes[i]) == h.codebook.M { st.NodesWithCode++ } } @@ -373,23 +393,20 @@ func (h *HNSW) EncodeAll(cb *PQCodebook) ([]uint64, [][]uint16, error) { return ids, codes, nil } -// LookupVectorByURL returns the persisted unit-normalized vector for the -// first passage whose url matches. Used by /find_similar?retriever=dense to -// skip the embed RPC — the source doc's vector is already in the graph. -// Linear scan; for 1M passages this is ~few ms, dominated by cache misses. -// Returns (nil, false) when the URL has no indexed passage. +// LookupVectorByURL returns the unit-normalized vector of the first live +// passage for url. Used by /find_similar?retriever=dense to skip the embed +// RPC. Returns (nil, false) when the URL has no live passage. func (h *HNSW) LookupVectorByURL(url string) ([]float32, bool) { h.mu.RLock() defer h.mu.RUnlock() - for i := range h.nodes { - if h.nodes[i].url == url { - // Copy: caller shouldn't be able to mutate graph internals. - cp := make([]float32, len(h.nodes[i].vec)) - copy(cp, h.nodes[i].vec) - return cp, true - } + ids := h.byURL[url] + if len(ids) == 0 { + return nil, false } - return nil, false + src := h.nodes[ids[0]].vec + cp := make([]float32, len(src)) + copy(cp, src) + return cp, true } // Add inserts a doc-level embedding without span info. Mirrors @@ -412,37 +429,61 @@ func (h *HNSW) codeFor(vec []float32) []uint16 { return code } -// MarkURLPassagesInvalid zeros out vec (and pq code, if present) for every -// node whose url matches. Returns the count zeroed. Dead nodes remain in -// the graph as link targets so neighbor adjacency lists stay consistent -// (searchLayer/Search both already skip nodes with empty vec — -// "zombie / partial-persisted" guard). Lets the crawler reclaim recall + -// memory on re-fetch instead of accumulating generations of stale chunks -// for the same URL. +// invalidateLocked turns node i into a zombie: vec and PQ code cleared, the +// node stays in place as a link target so adjacency lists stay consistent. +// Caller holds the write lock and owns the byURL entry. +func (h *HNSW) invalidateLocked(i int) { + h.nodes[i].vec = nil + if h.codes != nil && i < len(h.codes) { + h.codes[i] = nil + } + h.valid-- + h.dirty[int32(i)] = struct{}{} +} + +// relocateEntryLocked moves the entry point off a zombie to the highest-level +// live node; -1 when no live node remains. +func (h *HNSW) relocateEntryLocked() { + if h.entryPoint >= 0 && h.entryPoint < len(h.nodes) && len(h.nodes[h.entryPoint].vec) > 0 { + return + } + h.entryPoint = -1 + h.maxLevel = 0 + for i := range h.nodes { + if len(h.nodes[i].vec) > 0 && (h.entryPoint < 0 || h.nodes[i].level > h.maxLevel) { + h.entryPoint = i + h.maxLevel = h.nodes[i].level + } + } +} + +// MarkURLPassagesInvalid zombifies every live node for url (O(k) via the +// URL index). Returns the count zeroed. Lets the crawler reclaim recall + +// memory on re-fetch instead of accumulating generations of stale chunks. func (h *HNSW) MarkURLPassagesInvalid(url string) int { if url == "" { return 0 } h.mu.Lock() defer h.mu.Unlock() - n := 0 - for i := range h.nodes { - if h.nodes[i].url == url && len(h.nodes[i].vec) > 0 { - h.nodes[i].vec = nil - if h.codes != nil && i < len(h.codes) { - h.codes[i] = nil - } - n++ - } + ids := h.byURL[url] + if len(ids) == 0 { + return 0 } - return n + delete(h.byURL, url) + for _, id := range ids { + h.invalidateLocked(int(id)) + } + h.reclaimed.Add(uint64(len(ids))) + h.relocateEntryLocked() + return len(ids) } -// ReconcileURLs zeros out vec (and pq code) for every live node whose url -// fails the live predicate, in one pass under one lock acquisition. Used at -// load time to invalidate nodes whose docs were soft-deleted offline -// (purge-domain/purge-adult never touch the graph). Returns (invalidated, -// scanned). Idempotent: already-invalid nodes are skipped. +// ReconcileURLs zombifies every live node whose url fails the live +// predicate, in one pass under one lock acquisition. Used at load time to +// invalidate nodes whose docs were soft-deleted offline (purge-domain/ +// purge-adult never touch the graph). Returns (invalidated, scanned). +// Idempotent: already-invalid nodes are skipped. func (h *HNSW) ReconcileURLs(live func(url string) bool) (int, int) { if live == nil { return 0, 0 @@ -457,12 +498,13 @@ func (h *HNSW) ReconcileURLs(live func(url string) bool) (int, int) { if live(h.nodes[i].url) { continue } - h.nodes[i].vec = nil - if h.codes != nil && i < len(h.codes) { - h.codes[i] = nil - } + delete(h.byURL, h.nodes[i].url) + h.invalidateLocked(i) invalidated++ } + if invalidated > 0 { + h.relocateEntryLocked() + } return invalidated, len(h.nodes) } @@ -545,6 +587,8 @@ func (h *HNSW) addPassageLocked(url, title string, offset, length int, cp []floa level: level, neighbors: make([][]int, level+1), }) + h.byURL[url] = append(h.byURL[url], int32(newIdx)) + h.valid++ // when a codebook is loaded, encode the new vec inline and // keep h.codes parallel to h.nodes. The crawl-time PQ checkpoint // writes [lastN, len) of these to Pebble. @@ -560,9 +604,10 @@ func (h *HNSW) addPassageLocked(url, title string, offset, length int, cp []floa h.codes[newIdx] = code } - // First node: becomes the entry point trivially. - if newIdx == 0 { - h.entryPoint = 0 + // First node (or first live node after every other was invalidated): + // becomes the entry point trivially. + if h.entryPoint < 0 { + h.entryPoint = newIdx h.maxLevel = level return } @@ -750,11 +795,19 @@ func (h *HNSW) searchLayer(q []float32, pqTable []float32, entryPoints []int, ef results := &candMaxHeap{} heap.Init(results) + // Zombies are transit-only: expanded so traversal can cross them, never + // admitted to results (they would otherwise fill the ef window with + // +Inf entries — the failure mode when the entry point sits in an + // invalidated cluster). for _, ep := range entryPoints { + visited[ep] = struct{}{} + if h.zombieIdx(ep) { + heap.Push(cands, candEntry{idx: ep, dist: float32(math.Inf(1))}) + continue + } d := float32(h.distanceToNode(q, pqTable, ep)) heap.Push(cands, candEntry{idx: ep, dist: d}) heap.Push(results, candEntry{idx: ep, dist: d}) - visited[ep] = struct{}{} } for cands.Len() > 0 { @@ -784,6 +837,12 @@ func (h *HNSW) searchLayer(q []float32, pqTable []float32, entryPoints []int, ef continue } visited[nb] = struct{}{} + if h.zombieIdx(nb) { + if nb >= 0 && nb < len(h.nodes) { + heap.Push(cands, candEntry{idx: nb, dist: nearest.dist}) + } + continue + } d := float32(h.distanceToNode(q, pqTable, nb)) if results.Len() < ef { heap.Push(cands, candEntry{idx: nb, dist: d}) @@ -804,6 +863,11 @@ func (h *HNSW) searchLayer(q []float32, pqTable []float32, entryPoints []int, ef return out } +// zombieIdx reports whether idx is out of range or an invalidated node. +func (h *HNSW) zombieIdx(idx int) bool { + return idx < 0 || idx >= len(h.nodes) || len(h.nodes[idx].vec) == 0 +} + // addBackLink wires a back-edge from neighbor to newIdx at the given layer, // pruning neighbor's list if it overflows the per-layer cap. // @@ -819,6 +883,7 @@ func (h *HNSW) addBackLink(neighbor, newIdx, lvl int) { if lvl >= len(h.nodes[neighbor].neighbors) { return // neighbor doesn't participate at this layer } + h.dirty[int32(neighbor)] = struct{}{} mCap := h.M if lvl == 0 { mCap = h.Mmax0 diff --git a/internal/index/hnsw_compact.go b/internal/index/hnsw_compact.go index 9fcf13e..83eed7f 100644 --- a/internal/index/hnsw_compact.go +++ b/internal/index/hnsw_compact.go @@ -30,6 +30,7 @@ func (h *HNSW) Rebuild() *HNSW { fresh.efConstruction = h.efConstruction fresh.efSearch = h.efSearch fresh.levelMult = h.levelMult + fresh.slot = h.slot for i := range h.nodes { if len(h.nodes[i].vec) == 0 { @@ -115,9 +116,18 @@ func (h *HNSW) Compact() (removed int) { } } - // 3. Pick new entry point as the highest-level surviving node. + // 3. Rebuild the URL index and counters; every id changed, so the dirty + // set is meaningless until the caller's full persist rewrites the graph. h.nodes = newNodes h.codes = newCodes + h.byURL = make(map[string][]int32, len(h.byURL)) + for i := range h.nodes { + h.byURL[h.nodes[i].url] = append(h.byURL[h.nodes[i].url], int32(i)) + } + h.valid = len(h.nodes) + h.dirty = make(map[int32]struct{}) + + // 4. Pick new entry point as the highest-level surviving node. if len(h.nodes) == 0 { h.entryPoint = -1 h.maxLevel = 0 diff --git a/internal/index/hnsw_compact_test.go b/internal/index/hnsw_compact_test.go index ece45cd..c68f185 100644 --- a/internal/index/hnsw_compact_test.go +++ b/internal/index/hnsw_compact_test.go @@ -92,9 +92,10 @@ func TestHNSWZombieCompaction(t *testing.T) { } t.Logf("zombified %d of %d nodes", zombified, len(h.nodes)) + // Zombies are transit-only in searchLayer, so they cost nothing in recall. zombied := measure("with-zombies") - if zombied >= clean*0.9 { - t.Logf("note: zombies didn't drag recall enough on this seed (clean=%.3f, zombied=%.3f); test still validates compact()'s round-trip", clean, zombied) + if zombied < clean*0.95 { + t.Errorf("zombies dragged recall: clean=%.3f zombied=%.3f", clean, zombied) } removed := h.Compact() @@ -102,12 +103,11 @@ func TestHNSWZombieCompaction(t *testing.T) { t.Fatalf("compact removed %d, expected %d", removed, zombified) } + // Compact drops edges through removed nodes without re-linking survivors + // (Rebuild's job); recall may dip but must stay close to clean. compacted := measure("compacted") - // Compact doesn't restore edges to surviving nodes — it only removes - // dangling refs. Same recall as the zombied state (slightly higher is - // possible if traversal now skips fewer dead branches). - if compacted < zombied*0.95 { - t.Errorf("compaction regressed recall: zombied=%.3f compacted=%.3f", zombied, compacted) + if compacted < clean*0.9 { + t.Errorf("compaction regressed recall: clean=%.3f compacted=%.3f", clean, compacted) } // Rebuild does restore: fresh graph topology with full M-neighbor diff --git a/internal/index/hnsw_index_test.go b/internal/index/hnsw_index_test.go new file mode 100644 index 0000000..85a6941 --- /dev/null +++ b/internal/index/hnsw_index_test.go @@ -0,0 +1,315 @@ +package index + +import ( + "context" + "fmt" + "math/rand" + "path/filepath" + "testing" + + "github.com/pilot-protocol/cosift/internal/store" +) + +// checkURLIndex asserts the index invariant: a node is live iff its id is in +// byURL[url], and valid counts exactly the live nodes. +func checkURLIndex(t *testing.T, h *HNSW, where string) { + t.Helper() + h.mu.RLock() + defer h.mu.RUnlock() + live := 0 + indexed := map[int32]string{} + for url, ids := range h.byURL { + for _, id := range ids { + if _, dup := indexed[id]; dup { + t.Fatalf("%s: node %d indexed twice", where, id) + } + indexed[id] = url + } + } + for i := range h.nodes { + url, ok := indexed[int32(i)] + if len(h.nodes[i].vec) > 0 { + live++ + if !ok || url != h.nodes[i].url { + t.Fatalf("%s: live node %d (%s) missing from index (got %q)", where, i, h.nodes[i].url, url) + } + } else if ok { + t.Fatalf("%s: zombie node %d still indexed under %s", where, i, url) + } + } + if len(indexed) != live || h.valid != live { + t.Fatalf("%s: indexed=%d valid=%d live=%d", where, len(indexed), h.valid, live) + } + if h.entryPoint >= 0 && len(h.nodes[h.entryPoint].vec) == 0 { + t.Fatalf("%s: entry point %d is a zombie", where, h.entryPoint) + } + if h.entryPoint < 0 && live > 0 { + t.Fatalf("%s: no entry point with %d live nodes", where, live) + } +} + +func TestHNSWURLIndexInvariant(t *testing.T) { + h := buildTestHNSW(300, 8, 3, 5) + checkURLIndex(t, h, "build") + + batch := make([]PassageInput, 0, 12) + for i := 0; i < 12; i++ { + v := make([]float32, 8) + v[i%8] = 1 + batch = append(batch, PassageInput{URL: fmt.Sprintf("https://x/%d", i), Title: "regen", Offset: i, Length: 1, Vec: v}) + } + h.AddPassageBatch(batch) + checkURLIndex(t, h, "batch") + if len(h.byURL["https://x/1"]) != 2 { + t.Fatalf("second generation not indexed: %v", h.byURL["https://x/1"]) + } + + for i := 0; i < 40; i++ { + h.MarkURLPassagesInvalid(fmt.Sprintf("https://x/%d", i)) + } + checkURLIndex(t, h, "reclaim") + h.ReconcileURLs(func(u string) bool { return u != "https://x/100" && u != "https://x/200" }) + checkURLIndex(t, h, "reconcile") + + dir := filepath.Join(t.TempDir(), "pebble") + ps, err := store.OpenPebble(dir) + if err != nil { + t.Fatal(err) + } + defer ps.Close() + ctx := context.Background() + if err := h.Persist(ctx, ps); err != nil { + t.Fatal(err) + } + loaded, ok, err := LoadHNSW(ctx, ps) + if err != nil || !ok { + t.Fatalf("load: ok=%v err=%v", ok, err) + } + checkURLIndex(t, loaded, "load") + if loaded.valid != h.valid || len(loaded.byURL) != len(h.byURL) { + t.Fatalf("load counters: valid %d/%d urls %d/%d", loaded.valid, h.valid, len(loaded.byURL), len(h.byURL)) + } + + h.Compact() + checkURLIndex(t, h, "compact") + fresh := h.Rebuild() + checkURLIndex(t, fresh, "rebuild") + if fresh.Len() != h.Len() { + t.Fatalf("rebuild len %d != compact len %d", fresh.Len(), h.Len()) + } +} + +// Reclaim must touch only the URL's own nodes: the dirty set (every node +// written to) equals exactly the invalidated ids. +func TestHNSWReclaimTouchesOnlyURLNodes(t *testing.T) { + h := buildTestHNSW(5000, 8, 3, 5) + for i := 0; i < 3; i++ { + v := make([]float32, 8) + v[i] = 1 + h.AddPassage("https://multi", "m", i*10, 10, v) + } + want := map[int32]struct{}{} + for _, id := range h.byURL["https://multi"] { + want[id] = struct{}{} + } + h.dirty = make(map[int32]struct{}) + + if n := h.MarkURLPassagesInvalid("https://multi"); n != 3 { + t.Fatalf("reclaimed %d, want 3", n) + } + if len(h.dirty) != 3 { + t.Fatalf("reclaim touched %d nodes, want 3", len(h.dirty)) + } + for id := range want { + if _, ok := h.dirty[id]; !ok { + t.Fatalf("node %d not marked dirty", id) + } + } + if h.Reclaimed() != 3 || h.MarkURLPassagesInvalid("https://multi") != 0 { + t.Fatalf("reclaim counters: total=%d", h.Reclaimed()) + } + if _, ok := h.byURL["https://multi"]; ok { + t.Fatal("index entry survived reclaim") + } + checkURLIndex(t, h, "after") +} + +// Re-crawl: reclaim then add; search sees only the new generation and +// LookupVectorByURL never returns a zombie's empty vector. +func TestHNSWDuplicateGenerationReclaim(t *testing.T) { + h := NewHNSW(4) + h.AddPassage("https://doc", "v1", 0, 50, []float32{1, 0, 0, 0}) + h.AddPassage("https://doc", "v1", 50, 50, []float32{0, 1, 0, 0}) + h.Add("https://other", "o", []float32{0, 0, 0, 1}) + + if n := h.MarkURLPassagesInvalid("https://doc"); n != 2 { + t.Fatalf("reclaim %d, want 2", n) + } + if _, ok := h.LookupVectorByURL("https://doc"); ok { + t.Fatal("LookupVectorByURL returned a zombie") + } + h.AddPassage("https://doc", "v2", 0, 80, []float32{0, 0, 1, 0}) + v, ok := h.LookupVectorByURL("https://doc") + if !ok || v[2] < 0.99 { + t.Fatalf("lookup after re-add: ok=%v v=%v", ok, v) + } + hits := h.Search(context.Background(), []float32{1, 0, 0, 0}, 5) + for _, hit := range hits { + if hit.URL == "https://doc" && hit.Title != "v2" { + t.Fatalf("stale generation surfaced: %+v", hit) + } + } + st := h.PQStatus() + if st.NodesTotal != 4 || st.NodesValid != 2 { + t.Fatalf("status %+v", st) + } + checkURLIndex(t, h, "regen") +} + +func TestHNSWEntryPointRelocation(t *testing.T) { + h := buildTestHNSW(400, 8, 3, 5) + ep := h.entryPoint + epURL := h.nodes[ep].url + h.MarkURLPassagesInvalid(epURL) + if h.entryPoint == ep || h.zombieIdx(h.entryPoint) { + t.Fatalf("entry point not relocated: %d (was %d)", h.entryPoint, ep) + } + if h.nodes[h.entryPoint].level != h.maxLevel { + t.Fatalf("maxLevel %d != entry level %d", h.maxLevel, h.nodes[h.entryPoint].level) + } + q := make([]float32, 8) + q[0] = 1 + if hits := h.Search(context.Background(), q, 10); len(hits) != 10 { + t.Fatalf("search after relocation returned %d hits", len(hits)) + } + + inv, _ := h.ReconcileURLs(func(string) bool { return false }) + if inv != 399 || h.entryPoint != -1 { + t.Fatalf("reconcile all: inv=%d ep=%d", inv, h.entryPoint) + } + if hits := h.Search(context.Background(), q, 10); len(hits) != 0 { + t.Fatalf("all-zombie graph returned %d hits", len(hits)) + } + h.AddPassage("https://new", "n", 0, 1, q) + if h.entryPoint != 400 || h.maxLevel != h.nodes[400].level { + t.Fatalf("first live node did not become entry point: ep=%d", h.entryPoint) + } + if hits := h.Search(context.Background(), q, 10); len(hits) != 1 || hits[0].URL != "https://new" { + t.Fatalf("search after revival: %+v", hits) + } + checkURLIndex(t, h, "revived") +} + +// Regression for the 2026-08 production failure: an invalidated cluster +// around the entry point must not starve dense search. Zombies are +// traversed but never admitted to the ef window. +func TestHNSWSearchThroughZombieCluster(t *testing.T) { + const n, dim, k = 2000, 16, 10 + recall := func(h *HNSW) float64 { + rng := rand.New(rand.NewSource(77)) + total := 0.0 + const nq = 20 + for qi := 0; qi < nq; qi++ { + q := make([]float32, dim) + for j := range q { + q[j] = float32(rng.NormFloat64()) + } + gt := map[string]bool{} + for _, g := range h.BruteForceTopK(q, k) { + gt[g.URL] = true + } + hits := 0 + for _, a := range h.Search(context.Background(), q, k) { + if gt[a.URL] { + hits++ + } + } + total += float64(hits) / float64(len(gt)) + } + return total / nq + } + cluster := func(h *HNSW, hops int) map[int]struct{} { + set := map[int]struct{}{h.entryPoint: {}} + frontier := []int{h.entryPoint} + for hop := 0; hop < hops; hop++ { + var next []int + for _, i := range frontier { + for _, nb := range h.nodes[i].neighbors[0] { + if _, ok := set[nb]; !ok { + set[nb] = struct{}{} + next = append(next, nb) + } + } + } + frontier = next + } + return set + } + + // Variant A: the reclaim path (index-maintained, entry point relocates). + h := buildTestHNSW(n, dim, 3, 5) + clean := recall(h) + zombies := cluster(h, 2) + for i := range zombies { + h.MarkURLPassagesInvalid(h.nodes[i].url) + } + t.Logf("variant A: zombified %d nodes around the entry point", len(zombies)) + checkURLIndex(t, h, "cluster") + if r := recall(h); r < 0.9 || r < clean-0.05 { + t.Fatalf("variant A recall %.3f (clean %.3f)", r, clean) + } + + // Variant B: raw zombification with the entry point left in place — the + // on-disk shape after an offline purge on an old binary. + h = buildTestHNSW(n, dim, 3, 5) + zombies = cluster(h, 2) + for i := range zombies { + h.nodes[i].vec = nil + } + if !h.zombieIdx(h.entryPoint) { + t.Fatal("variant B: entry point should be a zombie") + } + hits := h.Search(context.Background(), make([]float32, dim), k) + for _, hit := range hits { + if _, z := zombies[int(h.byURL[hit.URL][0])]; z { + t.Fatalf("zombie surfaced: %+v", hit) + } + } + if r := recall(h); r < 0.9 || r < clean-0.05 { + t.Fatalf("variant B recall %.3f (clean %.3f)", r, clean) + } + + // Variant C: zombies as bridges — every layer-0 neighbour of the true + // top-k is invalidated, so the targets are reachable only through + // zombies. Without transit the ef window fills with finite results and + // the traversal never expands past the ring. + h = buildTestHNSW(n, dim, 3, 5) + rng := rand.New(rand.NewSource(77)) + ring := map[int]struct{}{} + for qi := 0; qi < 20; qi++ { + q := make([]float32, dim) + for j := range q { + q[j] = float32(rng.NormFloat64()) + } + targets := map[int]struct{}{} + for _, g := range h.BruteForceTopK(q, k) { + targets[int(h.byURL[g.URL][0])] = struct{}{} + } + for i := range targets { + for _, nb := range h.nodes[i].neighbors[0] { + if _, isTarget := targets[nb]; !isTarget { + ring[nb] = struct{}{} + } + } + } + } + for i := range ring { + if _, ok := h.byURL[h.nodes[i].url]; ok { + h.MarkURLPassagesInvalid(h.nodes[i].url) + } + } + t.Logf("variant C: zombified %d ring nodes", len(ring)) + if r := recall(h); r < 0.85 { + t.Fatalf("variant C recall %.3f (clean %.3f)", r, clean) + } +} diff --git a/internal/index/hnsw_persist.go b/internal/index/hnsw_persist.go index 837dd33..6d23665 100644 --- a/internal/index/hnsw_persist.go +++ b/internal/index/hnsw_persist.go @@ -27,11 +27,15 @@ // // Meta blob layout: // -// magic [4]byte = "HSW1" +// magic [4]byte = "HSW1" | "HSW2" // dim int32 // maxLevel int32 // entryPoint int32 // nodeCount int32 +// slot uint8 (HSW2 only; HSW1 implies slot 0x01) +// +// HSW1 is still emitted while the graph lives in slot 0x01 so a store that +// has never swapped stays readable by older binaries. package index @@ -41,12 +45,16 @@ import ( "fmt" "log" "math" + "sort" "time" "github.com/pilot-protocol/cosift/internal/store" ) -const hnswMetaMagic = "HSW1" +const ( + hnswMetaMagic = "HSW1" + hnswMetaMagicV2 = "HSW2" +) // persistWindowBytes bounds encoded blobs held in memory at once: a full // persist that materializes every blob first costs ~vec-bytes of extra heap @@ -56,50 +64,101 @@ var persistWindowBytes = 1 << 30 // persistFlushed is a test hook observing each flushed window (nil in prod). var persistFlushed func(nodes, bytes int) -// Persist serializes every node + meta into the PebbleStore. Safe to call -// during ongoing search (acquires RLock); does NOT acquire the write lock, -// so concurrent Add() during Persist will partially leak into the saved -// snapshot — callers expecting a clean snapshot should quiesce writes first. +// PersistProgress is reported by full persists after every flushed window. +type PersistProgress struct { + Written, Total int + Bytes int64 +} + +// Persist writes every node + meta into the graph's current slot. Prefer +// PersistSwap for full rewrites of a graph that already lives on disk. func (h *HNSW) Persist(ctx context.Context, ps *store.PebbleStore) error { return h.PersistFrom(ctx, ps, 0) } -// PersistFrom writes nodes[fromIdx:] in bounded windows, then meta. The -// crawl-time checkpoint goroutine uses this with fromIdx = last-persisted -// count, so each checkpoint touches only the newly-added nodes. Meta is -// always re-written so a reader can size the slice correctly. -// -// Caveat: existing nodes whose neighbor lists got new back-pointers since -// the last persist are NOT rewritten — those edges are lost until the next -// full-from-zero persist. Acceptable for crawl-time approximations; final -// shutdown persist always does fromIdx=0. +// PersistFrom writes nodes[fromIdx:] plus every node dirtied since the last +// persist (invalidated, or handed new back-links), then meta. The crawl-time +// checkpoint uses this with fromIdx = last-persisted count. Blocks while +// another persist or compact runs. func (h *HNSW) PersistFrom(ctx context.Context, ps *store.PebbleStore, fromIdx int) error { + h.persistMu.Lock() + defer h.persistMu.Unlock() + h.mu.RLock() + slot := h.slot + h.mu.RUnlock() + return h.persistNodes(ctx, ps, fromIdx, slot, nil) +} + +// TryPersistFrom is PersistFrom that returns (false, nil) instead of waiting +// when a persist or compact is already running. +func (h *HNSW) TryPersistFrom(ctx context.Context, ps *store.PebbleStore, fromIdx int) (bool, error) { + if !h.persistMu.TryLock() { + return false, nil + } + defer h.persistMu.Unlock() + h.mu.RLock() + slot := h.slot + h.mu.RUnlock() + return true, h.persistNodes(ctx, ps, fromIdx, slot, nil) +} + +// PersistSwap writes the whole graph into the inactive slot, then points meta +// at it. The previous slot stays loadable until the meta write; the caller +// clears it afterwards (store.OtherVectorSlot of the new Slot()). +func (h *HNSW) PersistSwap(ctx context.Context, ps *store.PebbleStore, progress func(PersistProgress)) error { + h.persistMu.Lock() + defer h.persistMu.Unlock() h.mu.RLock() - defer h.mu.RUnlock() - - if fromIdx >= len(h.nodes) { - // Even with no node writes, refresh meta so changes to maxLevel / - // entryPoint land on disk. - meta := encodeHNSWMeta(h.dim, h.maxLevel, h.entryPoint, len(h.nodes)) - return ps.PutVectorMeta(ctx, meta) - } - // The earlier order (meta then - // nodes) was unsafe — if the node batch failed, meta would point past - // actual data on disk and LoadHNSW would allocate slots for nodes that - // never landed, causing 'neighbors[-1]' panics during search. - // New order: meta ALWAYS lags or equals nodes-on-disk. Worst case after - // partial write: meta says N nodes, disk has N+M; the M extras are - // orphan but harmless (LoadHNSW caps at meta.nodeCount). - total := len(h.nodes) - fromIdx + old := h.slot + h.mu.RUnlock() + target := store.OtherVectorSlot(old) + if err := ps.ClearVectorSlot(ctx, target); err != nil { + return fmt.Errorf("clear target slot: %w", err) + } + if err := h.persistNodes(ctx, ps, 0, target, progress); err != nil { + return err + } + h.mu.Lock() + h.slot = target + h.mu.Unlock() + return nil +} + +// persistNodes is the shared write loop. Node blobs are encoded per window +// under the read lock and written outside it, so writers stall for one +// window at most. Meta is written last, from the same lock hold that +// observed the final node count, so it never points past written nodes. +// Caller holds persistMu. +func (h *HNSW) persistNodes(ctx context.Context, ps *store.PebbleStore, fromIdx int, slot byte, progress func(PersistProgress)) error { + h.mu.Lock() + dirty := h.dirty + h.dirty = make(map[int32]struct{}) + h.mu.Unlock() + dirtyIDs := make([]int32, 0, len(dirty)) + for id := range dirty { + if int(id) < fromIdx { + dirtyIDs = append(dirtyIDs, id) + } + } + sort.Slice(dirtyIDs, func(a, b int) bool { return dirtyIDs[a] < dirtyIDs[b] }) + restoreDirty := func() { + h.mu.Lock() + for id := range dirty { + h.dirty[id] = struct{}{} + } + h.mu.Unlock() + } + start := time.Now() window := make([]store.VectorNodeEntry, 0, 4096) windowBytes, written, flushes := 0, 0, 0 var bytesWritten int64 + total := 0 flush := func() error { if len(window) == 0 { return nil } - if err := ps.PutVectorNodesBatch(ctx, window); err != nil { + if err := ps.PutVectorNodesBatch(ctx, slot, window); err != nil { return fmt.Errorf("put vector nodes batch: %w", err) } written += len(window) @@ -115,26 +174,58 @@ func (h *HNSW) PersistFrom(ctx context.Context, ps *store.PebbleStore, fromIdx i log.Printf("hnsw persist: %d/%d nodes (%.1f GiB, %.0f nodes/s, eta %s)", written, total, float64(bytesWritten)/(1<<30), rate, eta) } + if progress != nil { + progress(PersistProgress{Written: written, Total: total, Bytes: bytesWritten}) + } clear(window) window = window[:0] windowBytes = 0 return nil } - for i := fromIdx; i < len(h.nodes); i++ { - blob := encodeHNSWNode(&h.nodes[i]) - window = append(window, store.VectorNodeEntry{ID: uint64(i), Blob: blob}) - windowBytes += len(blob) + 16 - if windowBytes >= persistWindowBytes { - if err := flush(); err != nil { - return err + + i, di := fromIdx, 0 + var meta []byte + for { + h.mu.RLock() + n := len(h.nodes) + total = max(n-fromIdx, 0) + len(dirtyIDs) + for windowBytes < persistWindowBytes { + var id int + switch { + case di < len(dirtyIDs): + id = int(dirtyIDs[di]) + di++ + case i < n: + id = i + i++ + default: + id = -1 } + if id < 0 { + break + } + if id >= n { + continue + } + blob := encodeHNSWNode(&h.nodes[id]) + window = append(window, store.VectorNodeEntry{ID: uint64(id), Blob: blob}) + windowBytes += len(blob) + 16 + } + done := di >= len(dirtyIDs) && i >= n + if done { + meta = encodeHNSWMeta(h.dim, h.maxLevel, h.entryPoint, n, slot) + } + h.mu.RUnlock() + if err := flush(); err != nil { + restoreDirty() + return err + } + if done { + break } } - if err := flush(); err != nil { - return err - } - meta := encodeHNSWMeta(h.dim, h.maxLevel, h.entryPoint, len(h.nodes)) if err := ps.PutVectorMeta(ctx, meta); err != nil { + restoreDirty() return fmt.Errorf("put vector meta: %w", err) } return nil @@ -199,6 +290,7 @@ func LoadHNSWProgress(ctx context.Context, ps *store.PebbleStore, progress func( h := NewHNSW(meta.dim) h.entryPoint = meta.entryPoint h.maxLevel = meta.maxLevel + h.slot = meta.slot h.nodes = make([]hnswNode, meta.nodeCount) // Corrupt blobs (bit rot, partial writes from prior crashes) leave the @@ -209,7 +301,7 @@ func LoadHNSWProgress(ctx context.Context, ps *store.PebbleStore, progress func( var loaded, processed uint64 var skipped int const logFirst = 5 - err = ps.IterateVectorNodes(ctx, func(nodeID uint64, blob []byte) bool { + err = ps.IterateVectorNodes(ctx, meta.slot, func(nodeID uint64, blob []byte) bool { if int(nodeID) >= len(h.nodes) { return true // out-of-bounds — skip silently } @@ -222,6 +314,10 @@ func LoadHNSWProgress(ctx context.Context, ps *store.PebbleStore, progress func( return true } h.nodes[nodeID] = *n + if len(n.vec) > 0 && n.url != "" { + h.byURL[n.url] = append(h.byURL[n.url], int32(nodeID)) + h.valid++ + } loaded++ processed++ if processed%loadCheckEvery == 0 { @@ -252,31 +348,43 @@ func LoadHNSWProgress(ctx context.Context, ps *store.PebbleStore, progress func( // hnswMetaDecoded is the in-memory shape of the meta blob. type hnswMetaDecoded struct { dim, maxLevel, entryPoint, nodeCount int + slot byte } -func encodeHNSWMeta(dim, maxLevel, entryPoint, nodeCount int) []byte { - buf := make([]byte, 4+4*4) +func encodeHNSWMeta(dim, maxLevel, entryPoint, nodeCount int, slot byte) []byte { + buf := make([]byte, 4+4*4, 4+4*4+1) copy(buf[0:4], hnswMetaMagic) binary.LittleEndian.PutUint32(buf[4:8], uint32(dim)) binary.LittleEndian.PutUint32(buf[8:12], uint32(maxLevel)) binary.LittleEndian.PutUint32(buf[12:16], uint32(entryPoint)) binary.LittleEndian.PutUint32(buf[16:20], uint32(nodeCount)) + if slot != store.VectorSlotA { + copy(buf[0:4], hnswMetaMagicV2) + buf = append(buf, slot) + } return buf } func decodeHNSWMeta(buf []byte) (hnswMetaDecoded, error) { - if len(buf) != 20 { - return hnswMetaDecoded{}, fmt.Errorf("meta blob: got %d bytes, want 20", len(buf)) - } - if string(buf[0:4]) != hnswMetaMagic { - return hnswMetaDecoded{}, fmt.Errorf("meta magic: got %q, want %q", buf[0:4], hnswMetaMagic) - } - return hnswMetaDecoded{ - dim: int(int32(binary.LittleEndian.Uint32(buf[4:8]))), - maxLevel: int(int32(binary.LittleEndian.Uint32(buf[8:12]))), - entryPoint: int(int32(binary.LittleEndian.Uint32(buf[12:16]))), - nodeCount: int(int32(binary.LittleEndian.Uint32(buf[16:20]))), - }, nil + var m hnswMetaDecoded + switch { + case len(buf) == 20 && string(buf[0:4]) == hnswMetaMagic: + m.slot = store.VectorSlotA + case len(buf) == 21 && string(buf[0:4]) == hnswMetaMagicV2: + m.slot = buf[20] + if m.slot != store.VectorSlotA && m.slot != store.VectorSlotB { + return hnswMetaDecoded{}, fmt.Errorf("meta slot: got %#x", m.slot) + } + case len(buf) != 20 && len(buf) != 21: + return hnswMetaDecoded{}, fmt.Errorf("meta blob: got %d bytes, want 20 or 21", len(buf)) + default: + return hnswMetaDecoded{}, fmt.Errorf("meta magic: got %q, want %q or %q", buf[0:4], hnswMetaMagic, hnswMetaMagicV2) + } + m.dim = int(int32(binary.LittleEndian.Uint32(buf[4:8]))) + m.maxLevel = int(int32(binary.LittleEndian.Uint32(buf[8:12]))) + m.entryPoint = int(int32(binary.LittleEndian.Uint32(buf[12:16]))) + m.nodeCount = int(int32(binary.LittleEndian.Uint32(buf[16:20]))) + return m, nil } func encodeHNSWNode(n *hnswNode) []byte { diff --git a/internal/index/hnsw_persist_test.go b/internal/index/hnsw_persist_test.go index ab16606..b7e6ba3 100644 --- a/internal/index/hnsw_persist_test.go +++ b/internal/index/hnsw_persist_test.go @@ -158,7 +158,7 @@ func TestLoadHNSWLegacyTruncatedZombieAccepted(t *testing.T) { if len(legacy) != 20 { t.Fatalf("legacy fixture size: got %d, want 20", len(legacy)) } - if err := ps.PutVectorNodesBatch(ctx, []store.VectorNodeEntry{ + if err := ps.PutVectorNodesBatch(ctx, store.VectorSlotA, []store.VectorNodeEntry{ {ID: legacyID, Blob: legacy}, }); err != nil { t.Fatalf("PutVectorNodesBatch: %v", err) @@ -335,7 +335,7 @@ func TestLoadHNSWSkipsCorruptNode(t *testing.T) { bad = append(bad, 0, 0, 0, 0) // length bad = append(bad, 0, 0, 0, 0) // level bad = append(bad, 99, 0, 0, 0) // dim = 99 — mismatch with meta dim=8 - if err := ps.PutVectorNodesBatch(ctx, []store.VectorNodeEntry{ + if err := ps.PutVectorNodesBatch(ctx, store.VectorSlotA, []store.VectorNodeEntry{ {ID: corruptID, Blob: bad}, }); err != nil { t.Fatalf("PutVectorNodesBatch: %v", err) diff --git a/internal/index/hnsw_swap_test.go b/internal/index/hnsw_swap_test.go new file mode 100644 index 0000000..c8ecaca --- /dev/null +++ b/internal/index/hnsw_swap_test.go @@ -0,0 +1,327 @@ +package index + +import ( + "context" + "fmt" + "math/rand" + "path/filepath" + "reflect" + "sync" + "testing" + + "github.com/pilot-protocol/cosift/internal/store" +) + +func openTestStore(t *testing.T) *store.PebbleStore { + t.Helper() + ps, err := store.OpenPebble(filepath.Join(t.TempDir(), "pebble")) + if err != nil { + t.Fatalf("OpenPebble: %v", err) + } + t.Cleanup(func() { ps.Close() }) + return ps +} + +func mustLoad(t *testing.T, ps *store.PebbleStore) *HNSW { + t.Helper() + g, ok, err := LoadHNSW(context.Background(), ps) + if err != nil || !ok { + t.Fatalf("load: ok=%v err=%v", ok, err) + } + return g +} + +// sameGraph compares every node's url/vec-liveness/neighbor lists. +func sameGraph(t *testing.T, want, got *HNSW, where string) { + t.Helper() + want.mu.RLock() + defer want.mu.RUnlock() + got.mu.RLock() + defer got.mu.RUnlock() + if len(want.nodes) != len(got.nodes) { + t.Fatalf("%s: len %d != %d", where, len(got.nodes), len(want.nodes)) + } + if want.entryPoint != got.entryPoint || want.maxLevel != got.maxLevel { + t.Fatalf("%s: meta ep %d/%d level %d/%d", where, got.entryPoint, want.entryPoint, got.maxLevel, want.maxLevel) + } + for i := range want.nodes { + w, g := &want.nodes[i], &got.nodes[i] + if w.url != g.url || (len(w.vec) > 0) != (len(g.vec) > 0) { + t.Fatalf("%s: node %d url/liveness differ", where, i) + } + for l := range w.neighbors { + wl, gl := w.neighbors[l], g.neighbors[l] + if len(wl) == 0 && len(gl) == 0 { + continue + } + if !reflect.DeepEqual(wl, gl) { + t.Fatalf("%s: node %d layer %d neighbors differ\n want %v\n got %v", where, i, l, wl, gl) + } + } + } +} + +func countSlot(t *testing.T, ps *store.PebbleStore, slot byte) int { + t.Helper() + n := 0 + if err := ps.IterateVectorNodes(context.Background(), slot, func(uint64, []byte) bool { n++; return true }); err != nil { + t.Fatal(err) + } + return n +} + +// Back-links added to already-persisted nodes land in the next incremental +// checkpoint. +func TestHNSWIncrementalPersistsBackLinks(t *testing.T) { + ps := openTestStore(t) + ctx := context.Background() + h := buildTestHNSW(200, 8, 3, 5) + if err := h.Persist(ctx, ps); err != nil { + t.Fatal(err) + } + rng := rand.New(rand.NewSource(9)) + for i := 200; i < 260; i++ { + v := make([]float32, 8) + for j := range v { + v[j] = float32(rng.NormFloat64()) + } + h.AddPassage(fmt.Sprintf("https://x/%d", i), "late", 0, 1, v) + } + if len(h.dirty) == 0 { + t.Fatal("inserts produced no dirty back-links") + } + if err := h.PersistFrom(ctx, ps, 200); err != nil { + t.Fatal(err) + } + if len(h.dirty) != 0 { + t.Fatalf("dirty set not drained: %d", len(h.dirty)) + } + sameGraph(t, h, mustLoad(t, ps), "after incremental") +} + +// Invalidations persist incrementally, including the fromIdx == len case +// that used to be a meta-only refresh. +func TestHNSWIncrementalPersistsInvalidations(t *testing.T) { + ps := openTestStore(t) + ctx := context.Background() + h := buildTestHNSW(100, 8, 3, 5) + if err := h.Persist(ctx, ps); err != nil { + t.Fatal(err) + } + h.MarkURLPassagesInvalid("https://x/7") + h.ReconcileURLs(func(u string) bool { return u != "https://x/8" }) + if err := h.PersistFrom(ctx, ps, 100); err != nil { + t.Fatal(err) + } + g := mustLoad(t, ps) + sameGraph(t, h, g, "after invalidation checkpoint") + if g.valid != 98 || len(g.byURL["https://x/7"]) != 0 || len(g.byURL["https://x/8"]) != 0 { + t.Fatalf("reloaded valid=%d byURL7=%v", g.valid, g.byURL["https://x/7"]) + } + checkURLIndex(t, g, "reload") +} + +func TestHNSWPersistErrorKeepsDirty(t *testing.T) { + ps := openTestStore(t) + ctx := context.Background() + h := buildTestHNSW(60, 8, 3, 5) + if err := h.Persist(ctx, ps); err != nil { + t.Fatal(err) + } + h.MarkURLPassagesInvalid("https://x/3") + dead := true + if err := h.PersistFrom(&failAfterFlagCtx{Context: ctx, fail: &dead}, ps, 60); err == nil { + t.Fatal("want persist error") + } + if _, ok := h.dirty[int32(3)]; !ok { + t.Fatal("dirty entry lost on persist failure") + } + if err := h.PersistFrom(ctx, ps, 60); err != nil { + t.Fatal(err) + } + g := mustLoad(t, ps) + if len(g.nodes[3].vec) != 0 { + t.Fatal("invalidation not persisted after retry") + } +} + +func TestHNSWPersistSwapRoundTrip(t *testing.T) { + ps := openTestStore(t) + ctx := context.Background() + h := buildTestHNSW(300, 8, 3, 5) + if err := h.Persist(ctx, ps); err != nil { + t.Fatal(err) + } + if meta, _, _ := ps.GetVectorMeta(ctx); len(meta) != 20 { + t.Fatalf("slot A meta must stay HSW1 (20 B), got %d", len(meta)) + } + h.MarkURLPassagesInvalid("https://x/1") + h.Compact() + + var last PersistProgress + if err := h.PersistSwap(ctx, ps, func(p PersistProgress) { last = p }); err != nil { + t.Fatal(err) + } + if h.Slot() != store.VectorSlotB || last.Written != 299 || last.Total != 299 { + t.Fatalf("slot=%#x progress=%+v", h.Slot(), last) + } + if meta, _, _ := ps.GetVectorMeta(ctx); len(meta) != 21 || meta[20] != store.VectorSlotB { + t.Fatalf("meta after swap: %v", meta) + } + if countSlot(t, ps, store.VectorSlotA) != 300 { + t.Fatal("old slot must survive until the caller clears it") + } + if err := ps.ClearVectorSlot(ctx, store.VectorSlotA); err != nil { + t.Fatal(err) + } + g := mustLoad(t, ps) + if g.Slot() != store.VectorSlotB { + t.Fatalf("loaded slot %#x", g.Slot()) + } + sameGraph(t, h, g, "after swap") + + // Incremental checkpoints now land in the new slot. + h.AddPassage("https://x/new", "n", 0, 1, []float32{1, 0, 0, 0, 0, 0, 0, 0}) + if err := h.PersistFrom(ctx, ps, 299); err != nil { + t.Fatal(err) + } + if countSlot(t, ps, store.VectorSlotB) != 300 || countSlot(t, ps, store.VectorSlotA) != 0 { + t.Fatalf("slots after incremental: A=%d B=%d", countSlot(t, ps, store.VectorSlotA), countSlot(t, ps, store.VectorSlotB)) + } + sameGraph(t, h, mustLoad(t, ps), "after incremental into B") + + // Swapping back lands in slot A with an HSW1 meta again. + if err := h.PersistSwap(ctx, ps, nil); err != nil { + t.Fatal(err) + } + if meta, _, _ := ps.GetVectorMeta(ctx); len(meta) != 20 || h.Slot() != store.VectorSlotA { + t.Fatalf("swap back: meta len %d slot %#x", len(meta), h.Slot()) + } + sameGraph(t, h, mustLoad(t, ps), "after swap back") +} + +// A swap that dies mid-write leaves meta on the old slot and the old graph +// loadable; the next swap clears the partial target and succeeds. +func TestHNSWPersistSwapFailureKeepsOldSlot(t *testing.T) { + ps := openTestStore(t) + ctx := context.Background() + h := buildTestHNSW(300, 16, 3, 5) + if err := h.Persist(ctx, ps); err != nil { + t.Fatal(err) + } + before := mustLoad(t, ps) + + h.MarkURLPassagesInvalid("https://x/2") + h.Compact() + savedWindow := persistWindowBytes + persistWindowBytes = 4 * 1024 + failNow := false + persistFlushed = func(int, int) { failNow = true } + defer func() { + persistWindowBytes = savedWindow + persistFlushed = nil + }() + if err := h.PersistSwap(&failAfterFlagCtx{Context: ctx, fail: &failNow}, ps, nil); err == nil { + t.Fatal("want swap error") + } + persistFlushed = nil + if h.Slot() != store.VectorSlotA { + t.Fatalf("slot flipped on failure: %#x", h.Slot()) + } + if countSlot(t, ps, store.VectorSlotB) == 0 { + t.Fatal("fixture: expected partial garbage in slot B") + } + sameGraph(t, before, mustLoad(t, ps), "old graph after failed swap") + + if err := h.PersistSwap(ctx, ps, nil); err != nil { + t.Fatal(err) + } + if countSlot(t, ps, store.VectorSlotB) != 299 { + t.Fatalf("slot B after retry: %d", countSlot(t, ps, store.VectorSlotB)) + } + sameGraph(t, h, mustLoad(t, ps), "after retry") +} + +func TestHNSWMetaFormats(t *testing.T) { + v1 := encodeHNSWMeta(8, 2, 5, 100, store.VectorSlotA) + if len(v1) != 20 || string(v1[:4]) != hnswMetaMagic { + t.Fatalf("slot A meta: %v", v1) + } + m, err := decodeHNSWMeta(v1) + if err != nil || m.slot != store.VectorSlotA || m.nodeCount != 100 { + t.Fatalf("decode v1: %+v %v", m, err) + } + v2 := encodeHNSWMeta(8, 2, 5, 100, store.VectorSlotB) + if len(v2) != 21 || string(v2[:4]) != hnswMetaMagicV2 { + t.Fatalf("slot B meta: %v", v2) + } + m, err = decodeHNSWMeta(v2) + if err != nil || m.slot != store.VectorSlotB || m.entryPoint != 5 { + t.Fatalf("decode v2: %+v %v", m, err) + } + bad := append([]byte{}, v2...) + bad[20] = 0x07 + if _, err := decodeHNSWMeta(bad); err == nil { + t.Fatal("bad slot accepted") + } + if _, err := decodeHNSWMeta(v2[:20]); err == nil { + t.Fatal("HSW2 magic with 20 bytes accepted") + } +} + +func TestHNSWTryPersistBusy(t *testing.T) { + ps := openTestStore(t) + h := buildTestHNSW(10, 8, 3, 5) + h.persistMu.Lock() + ok, err := h.TryPersistFrom(context.Background(), ps, 0) + h.persistMu.Unlock() + if ok || err != nil { + t.Fatalf("busy: ok=%v err=%v", ok, err) + } + ok, err = h.TryPersistFrom(context.Background(), ps, 0) + if !ok || err != nil { + t.Fatalf("idle: ok=%v err=%v", ok, err) + } + if mustLoad(t, ps).Len() != 10 { + t.Fatal("persist did not land") + } +} + +// Writers keep making progress while a windowed persist runs, and a final +// incremental checkpoint reconciles disk with memory. +func TestHNSWPersistConcurrentAdds(t *testing.T) { + ps := openTestStore(t) + ctx := context.Background() + h := buildTestHNSW(400, 8, 3, 5) + savedWindow := persistWindowBytes + persistWindowBytes = 8 * 1024 + defer func() { persistWindowBytes = savedWindow }() + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + rng := rand.New(rand.NewSource(21)) + for i := 0; i < 200; i++ { + v := make([]float32, 8) + for j := range v { + v[j] = float32(rng.NormFloat64()) + } + h.AddPassage(fmt.Sprintf("https://c/%d", i), "c", 0, 1, v) + if i%50 == 0 { + h.MarkURLPassagesInvalid(fmt.Sprintf("https://x/%d", i)) + } + } + }() + if err := h.Persist(ctx, ps); err != nil { + t.Fatal(err) + } + wg.Wait() + written := countSlot(t, ps, store.VectorSlotA) + if err := h.PersistFrom(ctx, ps, written); err != nil { + t.Fatal(err) + } + g := mustLoad(t, ps) + sameGraph(t, h, g, "after concurrent adds") + checkURLIndex(t, g, "reload") +} diff --git a/internal/store/pebble.go b/internal/store/pebble.go index 968bef8..7fd1747 100644 --- a/internal/store/pebble.go +++ b/internal/store/pebble.go @@ -153,6 +153,7 @@ func openPebble(path string, readOnly bool) (*PebbleStore, error) { cacheMB := envInt("COSIFT_PEBBLE_CACHE_MB", 128) memtableMB := envInt("COSIFT_PEBBLE_MEMTABLE_MB", 32) memtables := envInt("COSIFT_PEBBLE_MEMTABLES", 2) + compactions := envInt("COSIFT_PEBBLE_COMPACTIONS", 1) cache := pebble.NewCache(int64(cacheMB) << 20) defer cache.Unref() @@ -160,6 +161,7 @@ func openPebble(path string, readOnly bool) (*PebbleStore, error) { Cache: cache, MemTableSize: uint64(memtableMB) << 20, MemTableStopWritesThreshold: memtables + 2, + MaxConcurrentCompactions: func() int { return compactions }, ReadOnly: readOnly, } db, err := pebble.Open(path, opts) @@ -843,19 +845,37 @@ func docLenKey(docID int64) []byte { return k } -// Two sub-prefixes under the 'v' family: 0x00 for meta, 0x01 for nodes. -// Sorting on Pebble's byte ordering puts meta before nodes, so a startup -// iterator can read meta first and use it to size the node slice. +// Sub-prefixes under the 'v' family: 0x00 for meta, then one node slot per +// generation. A full persist writes the next generation into the inactive +// slot and only then points meta at it, so the previous graph stays +// loadable until the swap and the two key ranges never share tombstones. +const ( + VectorSlotA byte = 0x01 + VectorSlotB byte = 0x02 +) + +// OtherVectorSlot returns the inactive slot for the given active one. +func OtherVectorSlot(slot byte) byte { + if slot == VectorSlotA { + return VectorSlotB + } + return VectorSlotA +} + func vectorMetaKey() []byte { return []byte{famVector, 0x00, 'm', 'e', 't', 'a'} } -func vectorNodeKey(nodeID uint64) []byte { +func vectorNodeKey(slot byte, nodeID uint64) []byte { k := make([]byte, 1+1+8) k[0] = famVector - k[1] = 0x01 + k[1] = slot binary.BigEndian.PutUint64(k[2:], nodeID) return k } +func vectorSlotBounds(slot byte) (lo, hi []byte) { + return []byte{famVector, slot}, []byte{famVector, slot + 1} +} + // PutVectorMeta / GetVectorMeta — opaque blob holder under the 'v' family // for index-level metadata (entry point, max level, dim, node count). Format // is owned by the index package; store stays format-agnostic. @@ -886,11 +906,11 @@ func (p *PebbleStore) GetVectorMeta(ctx context.Context) ([]byte, bool, error) { } // PutVectorNode writes one HNSW-node blob under its ID. Caller-owned format. -func (p *PebbleStore) PutVectorNode(ctx context.Context, nodeID uint64, blob []byte) error { +func (p *PebbleStore) PutVectorNode(ctx context.Context, slot byte, nodeID uint64, blob []byte) error { if err := ctx.Err(); err != nil { return err } - return p.db.Set(vectorNodeKey(nodeID), blob, p.writeOpts) + return p.db.Set(vectorNodeKey(slot, nodeID), blob, p.writeOpts) } // VectorNodeEntry is one (id, blob) tuple for batched writes. @@ -908,7 +928,7 @@ type VectorNodeEntry struct { // // Used by index.HNSW.Persist for full snapshots and HNSW.PersistFrom for // incremental checkpoints. Iterb. -func (p *PebbleStore) PutVectorNodesBatch(ctx context.Context, entries []VectorNodeEntry) error { +func (p *PebbleStore) PutVectorNodesBatch(ctx context.Context, slot byte, entries []VectorNodeEntry) error { if err := ctx.Err(); err != nil { return err } @@ -939,7 +959,7 @@ func (p *PebbleStore) PutVectorNodesBatch(ctx context.Context, entries []VectorN batch = p.db.NewBatch() batchBytes = 0 } - if err := batch.Set(vectorNodeKey(e.ID), e.Blob, nil); err != nil { + if err := batch.Set(vectorNodeKey(slot, e.ID), e.Blob, nil); err != nil { return err } batchBytes += entryBytes @@ -1054,10 +1074,8 @@ func (p *PebbleStore) IteratePQCodes(ctx context.Context, fn func(nodeID uint64, return nil } -// ClearVectorFamily removes every persisted HNSW key (meta + nodes) in a -// single DeleteRange op. Used by `cosift hnsw-rebuild` before writing the -// freshly-reconstructed graph so leftover entries from the old graph -// can't shadow new ones at lower indices. +// ClearVectorFamily removes every persisted HNSW key (meta + all node slots) +// in a single DeleteRange op. func (p *PebbleStore) ClearVectorFamily(ctx context.Context) error { if err := ctx.Err(); err != nil { return err @@ -1067,6 +1085,33 @@ func (p *PebbleStore) ClearVectorFamily(ctx context.Context) error { return p.db.DeleteRange(lo, hi, p.writeOpts) } +// ClearVectorSlot removes every node blob in one slot and compacts the range +// synchronously so a later persist into it never lands on top of tombstones. +func (p *PebbleStore) ClearVectorSlot(ctx context.Context, slot byte) error { + if err := ctx.Err(); err != nil { + return err + } + lo, hi := vectorSlotBounds(slot) + if err := p.db.DeleteRange(lo, hi, p.writeOpts); err != nil { + return err + } + return p.db.Compact(lo, hi, true) +} + +// VectorSlotEmpty reports whether a node slot holds no entries. +func (p *PebbleStore) VectorSlotEmpty(ctx context.Context, slot byte) (bool, error) { + if err := ctx.Err(); err != nil { + return false, err + } + lo, hi := vectorSlotBounds(slot) + it, err := p.db.NewIter(&pebble.IterOptions{LowerBound: lo, UpperBound: hi}) + if err != nil { + return false, err + } + defer it.Close() + return !it.First(), nil +} + // DB returns the underlying Pebble handle for operations that need // direct access (e.g. the offline pebble-compact subcommand calling // db.Compact() to collapse tombstones after a frontier wipe). @@ -1112,14 +1157,14 @@ func (p *PebbleStore) ClearPQFamily(ctx context.Context) error { return p.db.DeleteRange(lo, hi, p.writeOpts) } -// IterateVectorNodes scans every persisted HNSW node in ascending ID order, -// invoking fn(nodeID, blob) for each. Returning false from fn stops the scan. -func (p *PebbleStore) IterateVectorNodes(ctx context.Context, fn func(nodeID uint64, blob []byte) bool) error { +// IterateVectorNodes scans every persisted HNSW node in one slot in ascending +// ID order, invoking fn(nodeID, blob) for each. Returning false from fn stops +// the scan. +func (p *PebbleStore) IterateVectorNodes(ctx context.Context, slot byte, fn func(nodeID uint64, blob []byte) bool) error { if err := ctx.Err(); err != nil { return err } - lo := []byte{famVector, 0x01} - hi := []byte{famVector, 0x02} + lo, hi := vectorSlotBounds(slot) it, err := p.db.NewIter(&pebble.IterOptions{LowerBound: lo, UpperBound: hi}) if err != nil { return err diff --git a/internal/store/pebble_uncovered_test.go b/internal/store/pebble_uncovered_test.go index a45db76..5e764df 100644 --- a/internal/store/pebble_uncovered_test.go +++ b/internal/store/pebble_uncovered_test.go @@ -72,13 +72,13 @@ func TestPebbleVectorNodeRoundtrip(t *testing.T) { p := newPebbleStore(t) ctx := context.Background() - if err := p.PutVectorNode(ctx, 42, []byte("nodepayload")); err != nil { + if err := p.PutVectorNode(ctx, VectorSlotA, 42, []byte("nodepayload")); err != nil { t.Fatalf("PutVectorNode: %v", err) } // Iterate and confirm we see it. found := false - if err := p.IterateVectorNodes(ctx, func(id uint64, blob []byte) bool { + if err := p.IterateVectorNodes(ctx, VectorSlotA, func(id uint64, blob []byte) bool { if id == 42 && string(blob) == "nodepayload" { found = true } @@ -96,7 +96,7 @@ func TestPebbleVectorNodesBatch(t *testing.T) { ctx := context.Background() // Empty input is a no-op. - if err := p.PutVectorNodesBatch(ctx, nil); err != nil { + if err := p.PutVectorNodesBatch(ctx, VectorSlotA, nil); err != nil { t.Errorf("empty batch: %v", err) } @@ -105,12 +105,12 @@ func TestPebbleVectorNodesBatch(t *testing.T) { {ID: 2, Blob: []byte("two")}, {ID: 3, Blob: []byte("three")}, } - if err := p.PutVectorNodesBatch(ctx, entries); err != nil { + if err := p.PutVectorNodesBatch(ctx, VectorSlotA, entries); err != nil { t.Fatalf("PutVectorNodesBatch: %v", err) } seen := map[uint64]string{} - _ = p.IterateVectorNodes(ctx, func(id uint64, blob []byte) bool { + _ = p.IterateVectorNodes(ctx, VectorSlotA, func(id uint64, blob []byte) bool { seen[id] = string(blob) return true }) @@ -119,11 +119,45 @@ func TestPebbleVectorNodesBatch(t *testing.T) { } } +// Slots are disjoint key ranges: clearing one leaves the other intact. +func TestPebbleVectorSlots(t *testing.T) { + p := newPebbleStore(t) + ctx := context.Background() + _ = p.PutVectorNode(ctx, VectorSlotA, 1, []byte("a1")) + _ = p.PutVectorNode(ctx, VectorSlotB, 1, []byte("b1")) + + if OtherVectorSlot(VectorSlotA) != VectorSlotB || OtherVectorSlot(VectorSlotB) != VectorSlotA { + t.Fatal("OtherVectorSlot mapping") + } + for _, slot := range []byte{VectorSlotA, VectorSlotB} { + if empty, err := p.VectorSlotEmpty(ctx, slot); err != nil || empty { + t.Fatalf("slot %d: empty=%v err=%v", slot, empty, err) + } + } + if err := p.ClearVectorSlot(ctx, VectorSlotA); err != nil { + t.Fatalf("ClearVectorSlot: %v", err) + } + if empty, _ := p.VectorSlotEmpty(ctx, VectorSlotA); !empty { + t.Error("slot A should be empty") + } + got := "" + _ = p.IterateVectorNodes(ctx, VectorSlotB, func(_ uint64, blob []byte) bool { + got = string(blob) + return true + }) + if got != "b1" { + t.Errorf("slot B disturbed: %q", got) + } + if err := p.ClearVectorSlot(ctx, VectorSlotA); err != nil { + t.Fatalf("ClearVectorSlot on empty slot: %v", err) + } +} + func TestPebbleClearVectorFamily(t *testing.T) { p := newPebbleStore(t) ctx := context.Background() _ = p.PutVectorMeta(ctx, []byte("meta")) - _ = p.PutVectorNode(ctx, 1, []byte("one")) + _ = p.PutVectorNode(ctx, VectorSlotA, 1, []byte("one")) if err := p.ClearVectorFamily(ctx); err != nil { t.Fatalf("ClearVectorFamily: %v", err) @@ -132,7 +166,7 @@ func TestPebbleClearVectorFamily(t *testing.T) { t.Errorf("VectorMeta should be cleared") } count := 0 - _ = p.IterateVectorNodes(ctx, func(_ uint64, _ []byte) bool { + _ = p.IterateVectorNodes(ctx, VectorSlotA, func(_ uint64, _ []byte) bool { count++ return true }) From 9961582a3930481eaf036659c22553ab628393b5 Mon Sep 17 00:00:00 2001 From: Andrei Marinescu Date: Wed, 2 Sep 2026 10:59:35 +0300 Subject: [PATCH 2/3] serve: async hnsw-compact job with /stats progress; incremental shutdown checkpoint; reclaim on by default /admin/hnsw-compact now returns 202 and runs CompactPersist in the background (409 while in flight, ?wait=1 blocks for the result); state, phase, persist progress and ETA are published as /stats.hnsw_compact. The old slot is cleared only after the swap. /admin/checkpoint and /admin/frontier-clear lift the 60 s WriteTimeout like hnsw-compact already did. The crawler checkpoint uses TryPersistFrom (skips a tick while a full persist runs) and also fires on dirty-only changes; shutdown is one more incremental checkpoint instead of a full persist that never completed at production scale. A stale slot left by a crash mid-swap is cleared after load. COSIFT_ZOMBIE_RECLAIM defaults to on (0/false/off disables); reclaim and zombie counts are exposed in /stats and /metrics. COSIFT_PEBBLE_COMPACTIONS sets Pebble MaxConcurrentCompactions (default 1). --- cmd/cosift/dense_reconcile_test.go | 71 +++++++- cmd/cosift/hnsw_rebuild.go | 25 +-- cmd/cosift/serve_admin.go | 221 ++++++++++++++++-------- cmd/cosift/serve_crawl.go | 3 + cmd/cosift/serve_setup.go | 112 ++++++------ cmd/cosift/serve_stats.go | 21 +++ internal/crawler/crawler.go | 16 +- internal/crawler/store_iface.go | 12 ++ internal/crawler/wet.go | 6 +- internal/crawler/zombie_reclaim_test.go | 15 ++ internal/index/hnsw_compact.go | 95 +++++++++- internal/index/hnsw_persist.go | 4 + 12 files changed, 456 insertions(+), 145 deletions(-) create mode 100644 internal/crawler/zombie_reclaim_test.go diff --git a/cmd/cosift/dense_reconcile_test.go b/cmd/cosift/dense_reconcile_test.go index 60275b8..24208d5 100644 --- a/cmd/cosift/dense_reconcile_test.go +++ b/cmd/cosift/dense_reconcile_test.go @@ -7,6 +7,10 @@ import ( "net/http/httptest" "strings" "testing" + "time" + + "github.com/pilot-protocol/cosift/internal/index" + "github.com/pilot-protocol/cosift/internal/store" ) // divergeFixture soft-deletes the docs behind the given URLs, leaving their @@ -238,7 +242,11 @@ func TestHNSWCompactPersistHardening(t *testing.T) { compact := func(q string, ctx context.Context) map[string]any { t.Helper() - req := httptest.NewRequest(http.MethodPost, "/admin/hnsw-compact"+q, nil) + sep := "?" + if q != "" { + sep = "&" + } + req := httptest.NewRequest(http.MethodPost, "/admin/hnsw-compact"+q+sep+"wait=1", nil) if ctx != nil { req = req.WithContext(ctx) } @@ -280,6 +288,67 @@ func TestHNSWCompactPersistHardening(t *testing.T) { if resp["removed"].(float64) != 1 || resp["persisted"] != false { t.Fatalf("skip_persist changed behavior: %v", resp) } + // Each persisting run swapped slots; the graph reloads from the active one + // and the old slot is empty. + if f.hnsw.Slot() != store.VectorSlotA { + t.Fatalf("two swaps should land back in slot A, got %#x", f.hnsw.Slot()) + } + if empty, _ := f.ps.VectorSlotEmpty(context.Background(), store.VectorSlotB); !empty { + t.Fatal("old slot not cleared after swap") + } + g, ok, err := index.LoadHNSW(context.Background(), f.ps) + if err != nil || !ok || g.Len() != f.hnsw.Len()+1 { + t.Fatalf("reload after compact: ok=%v err=%v len=%d", ok, err, g.Len()) + } +} + +// The async path: 202 on start, 409 while running, progress + result in +// /stats.hnsw_compact once done. +func TestHNSWCompactAsyncJob(t *testing.T) { + f := populatedPebbleStore(t) + srv := f.makeServer(nil) + f.hnsw.MarkURLPassagesInvalid(f.docs[5]) + post := func(q string) (int, map[string]any) { + req := httptest.NewRequest(http.MethodPost, "/admin/hnsw-compact"+q, nil) + rec := httptest.NewRecorder() + srv.handleHNSWCompact(rec, req) + var body map[string]any + _ = json.Unmarshal(rec.Body.Bytes(), &body) + return rec.Code, body + } + code, body := post("") + if code != http.StatusAccepted || body["status"] != "started" { + t.Fatalf("start: %d %v", code, body) + } + for i := 0; i < 200; i++ { + if srv.compact.snapshot()["state"] != "running" { + break + } + if c, _ := post(""); c != http.StatusConflict && c != http.StatusAccepted { + t.Fatalf("second POST while running: %d", c) + } + time.Sleep(10 * time.Millisecond) + } + srv.bgJobs.Wait() + snap := srv.compact.snapshot() + if snap["state"] != "done" || snap["phase"] != "done" || snap["persisted"] != true { + t.Fatalf("snapshot after run: %v", snap) + } + if snap["removed"].(int) != 1 { + t.Fatalf("removed: %v", snap["removed"]) + } + raw, err := srv.buildStatsBody(context.Background()) + if err != nil { + t.Fatal(err) + } + var stats map[string]any + _ = json.Unmarshal(raw, &stats) + if _, ok := stats["hnsw_compact"]; !ok { + t.Fatalf("/stats lacks hnsw_compact: %v", stats) + } + if stats["hnsw_reclaimed_total"].(float64) < 1 { + t.Fatalf("hnsw_reclaimed_total: %v", stats["hnsw_reclaimed_total"]) + } } // /answer under divergence (defaults to hybrid when dense is ready) counts diff --git a/cmd/cosift/hnsw_rebuild.go b/cmd/cosift/hnsw_rebuild.go index 5aa7788..b769e52 100644 --- a/cmd/cosift/hnsw_rebuild.go +++ b/cmd/cosift/hnsw_rebuild.go @@ -13,9 +13,9 @@ import ( // runHNSWRebuild reconstructs the HNSW graph in a pebble dir from valid // (vec != nil) nodes only. Removes zombies, recovers full M-neighbor -// connectivity, persists the fresh graph back into the same dir under a -// cleared 'v' family. Invalidates 'q' (PQ) — operators must re-train PQ -// after rebuild via /admin/pq-train. +// connectivity, persists the fresh graph into the inactive node slot and +// swaps. Invalidates 'q' (PQ) — operators must re-train PQ after rebuild via +// /admin/pq-train. // // Use this against a Pebble dir whose serve has been stopped (Pebble locks // the dir exclusively). Pair with /admin/checkpoint to take a consistent @@ -76,20 +76,21 @@ func runHNSWRebuild(ctx context.Context, cfg *config.Config, args []string) erro return nil } + persistT0 := time.Now() + old := fresh.Slot() + if err := fresh.PersistSwap(ctx, ps, nil); err != nil { + return fmt.Errorf("persist: %w", err) + } + fmt.Printf("hnsw-rebuild: persisted into slot %#x in %s\n", fresh.Slot(), time.Since(persistT0).Round(time.Second)) + clearT0 := time.Now() - if err := ps.ClearVectorFamily(ctx); err != nil { - return fmt.Errorf("clear vector family: %w", err) + if err := ps.ClearVectorSlot(ctx, old); err != nil { + return fmt.Errorf("clear old slot: %w", err) } if err := ps.ClearPQFamily(ctx); err != nil { return fmt.Errorf("clear PQ family: %w", err) } - fmt.Printf("hnsw-rebuild: cleared old v + q families in %s\n", time.Since(clearT0).Round(time.Millisecond)) - - persistT0 := time.Now() - if err := fresh.Persist(ctx, ps); err != nil { - return fmt.Errorf("persist: %w", err) - } - fmt.Printf("hnsw-rebuild: persisted in %s\n", time.Since(persistT0).Round(time.Second)) + fmt.Printf("hnsw-rebuild: cleared old slot + q family in %s\n", time.Since(clearT0).Round(time.Millisecond)) fmt.Printf("hnsw-rebuild: DONE. nodes %d → %d (PQ codes cleared; run /admin/pq-train to restore PQ)\n", g.Len(), fresh.Len()) diff --git a/cmd/cosift/serve_admin.go b/cmd/cosift/serve_admin.go index 8265b82..e98251b 100644 --- a/cmd/cosift/serve_admin.go +++ b/cmd/cosift/serve_admin.go @@ -91,6 +91,9 @@ func (s *pebbleHTTP) handleCheckpoint(w http.ResponseWriter, r *http.Request) { writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") return } + if rc := http.NewResponseController(w); rc != nil { + _ = rc.SetWriteDeadline(time.Time{}) + } base := os.Getenv("COSIFT_CHECKPOINT_DIR") if base == "" { base = "/tmp" @@ -335,90 +338,170 @@ func (s *pebbleHTTP) handleEvalQuick(w http.ResponseWriter, r *http.Request) { }) } -// handleHNSWCompact runs HNSW.Compact() in-place, then clears the persisted -// 'v' family and writes a fresh full snapshot so disk matches the compacted -// in-memory graph. Cheaper than the offline hnsw-rebuild subcommand: Compact -// keeps the existing topology among surviving nodes (O(N + edges)), whereas -// Rebuild re-inserts every node via HNSW search (multiple minutes per million -// passages). Operators run this when stats.zombie_nodes climbs above ~30% of -// nodes_total. -// -// Synchronous; holds the HNSW write lock during the compact step and the -// read lock during the persist step. Dense retrieval and AddPassage calls -// queue for the duration. The server-wide WriteTimeout is disabled here via -// ResponseController because compacting a multi-million-node graph routinely -// runs past 60s. Returns counters so operators can confirm progress. +// compactJob is the single-slot state of the async /admin/hnsw-compact run, +// mirrored into /stats.hnsw_compact. +type compactJob struct { + mu sync.Mutex + running bool + started time.Time + finished time.Time + progress index.CompactProgress + result index.CompactResult + err error + done chan struct{} +} + +func (j *compactJob) snapshot() map[string]any { + j.mu.Lock() + defer j.mu.Unlock() + return j.snapshotLocked() +} + +func (j *compactJob) snapshotLocked() map[string]any { + m := map[string]any{"state": "idle"} + if j.started.IsZero() { + return m + } + switch { + case j.running: + m["state"] = "running" + case j.err != nil: + m["state"] = "error" + m["error"] = j.err.Error() + default: + m["state"] = "done" + } + p := j.progress + m["phase"] = p.Phase + m["nodes_before"] = p.NodesBefore + m["nodes_after"] = p.NodesAfter + m["removed"] = p.Removed + m["started_at"] = j.started.UTC().Format(time.RFC3339) + end := j.finished + if j.running { + end = time.Now() + } else { + m["finished_at"] = end.UTC().Format(time.RFC3339) + } + elapsed := end.Sub(j.started).Seconds() + m["elapsed_s"] = elapsed + if p.Total > 0 { + m["persist_written"] = p.Written + m["persist_total"] = p.Total + m["persist_pct"] = 100 * float64(p.Written) / float64(p.Total) + if j.running && p.Phase == "persist" && p.Written > 0 && elapsed > 0 { + m["eta_s"] = float64(p.Total-p.Written) / (float64(p.Written) / elapsed) + } + } + if !j.running { + m["persisted"] = j.result.Persisted + m["compact_ms"] = j.result.CompactDur.Milliseconds() + m["persist_ms"] = j.result.PersistDur.Milliseconds() + } + return m +} + +// resultJSON is the completion payload (also returned by ?wait=1). +func (j *compactJob) resultJSON() (map[string]any, int) { + j.mu.Lock() + defer j.mu.Unlock() + r := j.result + resp := map[string]any{ + "nodes_before": r.NodesBefore, + "nodes_after": r.NodesAfter, + "removed": r.Removed, + "compact_ms": r.CompactDur.Milliseconds(), + "persisted": r.Persisted, + } + if r.Forced { + resp["forced"] = true + } + if r.Persisted { + resp["persist_ms"] = r.PersistDur.Milliseconds() + } + if j.err != nil { + resp["persist_error"] = j.err.Error() + return resp, http.StatusInternalServerError + } + return resp, http.StatusOK +} + +// handleHNSWCompact starts HNSW.CompactPersist as a background job: compact +// in place, rewrite the graph into the inactive on-disk slot, swap, clear the +// old slot. Returns 202 immediately (409 while a run is in flight); progress +// lives in /stats.hnsw_compact. ?wait=1 blocks for the result instead — +// the WriteTimeout is lifted for that case. Options: skip_persist=1, +// force_persist=1 (re-persist even when nothing was removed — the retry +// path after an interrupted run). PQ codes are cleared with the old slot; +// operators re-run /admin/pq-train if PQ was in use. func (s *pebbleHTTP) handleHNSWCompact(w http.ResponseWriter, r *http.Request) { if !peerTokenOK(r, s.cluster.PeerAuthToken) { writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") return } - if s.hnsw() == nil { + g := s.hnsw() + if g == nil { writeProblem(w, http.StatusNotImplemented, "hnsw-compact requires a loaded HNSW graph") return } - if rc := http.NewResponseController(w); rc != nil { - _ = rc.SetWriteDeadline(time.Time{}) - } - skipPersist := r.URL.Query().Get("skip_persist") == "1" - forcePersist := r.URL.Query().Get("force_persist") == "1" + q := r.URL.Query() + skipPersist := q.Get("skip_persist") == "1" + forcePersist := q.Get("force_persist") == "1" + wait := q.Get("wait") == "1" - before := s.hnsw().Len() - t0 := time.Now() - removed := s.hnsw().Compact() - compactDur := time.Since(t0) - after := s.hnsw().Len() - - resp := map[string]any{ - "nodes_before": before, - "nodes_after": after, - "removed": removed, - "compact_ms": compactDur.Milliseconds(), - "persisted": false, - } - - // force_persist=1 re-runs the wipe+persist even when this compact removed - // nothing — the retry path after a failed or interrupted persist, which - // otherwise leaves the disk graph partial with no way to repair it - // in-process (a second compact finds removed==0 and returns here). - if skipPersist || (removed == 0 && !forcePersist) { - writeJSON(w, http.StatusOK, resp) + j := &s.compact + j.mu.Lock() + if j.running { + snap := j.snapshotLocked() + j.mu.Unlock() + writeJSON(w, http.StatusConflict, snap) return } - if forcePersist { - resp["forced"] = true - } + j.running = true + j.started = time.Now() + j.finished = time.Time{} + j.progress = index.CompactProgress{Phase: "compact", NodesBefore: g.Len()} + j.result = index.CompactResult{} + j.err = nil + done := make(chan struct{}) + j.done = done + j.mu.Unlock() - // Compact remapped node indices; the persisted 'v' family now points at - // stale slots. Wipe and full-rewrite. PQ codes follow node indices too, - // so clear 'q' as well — operators must re-run /admin/pq-train if PQ was - // in use. - persistT0 := time.Now() - // Deliberately NOT r.Context(): a dropped client connection mid-persist - // would cancel the wipe+rewrite and strand a partial disk graph. - ctx := context.Background() - if err := s.store.ClearVectorFamily(ctx); err != nil { - resp["persist_error"] = "clear vector family: " + err.Error() - writeJSON(w, http.StatusInternalServerError, resp) - return - } - if err := s.store.ClearPQFamily(ctx); err != nil { - resp["persist_error"] = "clear pq family: " + err.Error() - writeJSON(w, http.StatusInternalServerError, resp) + s.bgJobs.Add(1) + go func() { + defer s.bgJobs.Done() + // Deliberately not r.Context(): the run must outlive the request. + res, err := g.CompactPersist(context.Background(), s.store, skipPersist, forcePersist, func(p index.CompactProgress) { + j.mu.Lock() + j.progress = p + j.mu.Unlock() + }) + j.mu.Lock() + j.result, j.err = res, err + j.running = false + j.finished = time.Now() + j.mu.Unlock() + close(done) + if err != nil { + log.Printf("hnsw-compact: FAILED after removing %d nodes: %v", res.Removed, err) + return + } + log.Printf("hnsw-compact: removed=%d (%.1f%% zombies) compact=%s persist=%s persisted=%v nodes %d→%d slot=%#x", + res.Removed, 100*float64(res.Removed)/float64(max(res.NodesBefore, 1)), + res.CompactDur.Round(time.Millisecond), res.PersistDur.Round(time.Millisecond), + res.Persisted, res.NodesBefore, res.NodesAfter, g.Slot()) + }() + + if !wait { + writeJSON(w, http.StatusAccepted, map[string]any{"status": "started", "watch": "/stats hnsw_compact"}) return } - if err := s.hnsw().Persist(ctx, s.store); err != nil { - resp["persist_error"] = "persist: " + err.Error() - writeJSON(w, http.StatusInternalServerError, resp) - return + if rc := http.NewResponseController(w); rc != nil { + _ = rc.SetWriteDeadline(time.Time{}) } - resp["persisted"] = true - resp["persist_ms"] = time.Since(persistT0).Milliseconds() - log.Printf("hnsw-compact: removed=%d (%.1f%% zombies) compact=%s persist=%s nodes %d→%d", - removed, 100*float64(removed)/float64(before), - compactDur.Round(time.Millisecond), time.Since(persistT0).Round(time.Millisecond), - before, after) - writeJSON(w, http.StatusOK, resp) + <-done + resp, code := j.resultJSON() + writeJSON(w, code, resp) } // responseRecorder captures an http.Handler's output for in-process diff --git a/cmd/cosift/serve_crawl.go b/cmd/cosift/serve_crawl.go index c5ef940..a8efc13 100644 --- a/cmd/cosift/serve_crawl.go +++ b/cmd/cosift/serve_crawl.go @@ -98,6 +98,9 @@ func (s *pebbleHTTP) handleFrontierClear(w http.ResponseWriter, r *http.Request) writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") return } + if rc := http.NewResponseController(w); rc != nil { + _ = rc.SetWriteDeadline(time.Time{}) + } if err := s.store.ClearFrontier(r.Context()); err != nil { writeProblem(w, http.StatusInternalServerError, err.Error()) return diff --git a/cmd/cosift/serve_setup.go b/cmd/cosift/serve_setup.go index ad01f3f..61d2165 100644 --- a/cmd/cosift/serve_setup.go +++ b/cmd/cosift/serve_setup.go @@ -658,6 +658,7 @@ func runPebbleServe(ctx context.Context, cfg *config.Config, args []string) erro servErr := httpSrv.ListenAndServe() bgWG.Wait() // loader goroutine (and its crawler-start decision) done crawlWG.Wait() // crawler final persist before the deferred ps.Close() + srv.bgJobs.Wait() if servErr != nil && servErr != http.ErrServerClosed { return servErr } @@ -838,8 +839,19 @@ func (s *pebbleHTTP) loadHNSWInto(ctx context.Context, ps *store.PebbleStore, ve s.hnswAt.Store(g) s.hnswLoaded.Store(uint64(g.Len())) s.hnswLoadState.Store(2) // ready - log.Printf("pebble-serve: HNSW graph loaded into memory: %d nodes, dim=%d (%.0fs)", - g.Len(), vectorDim, time.Since(start).Seconds()) + log.Printf("pebble-serve: HNSW graph loaded into memory: %d nodes, dim=%d, slot=%#x (%.0fs)", + g.Len(), vectorDim, g.Slot(), time.Since(start).Seconds()) + // A crash between a slot swap and its cleanup leaves the previous + // generation on disk; reclaim it before the crawler starts persisting. + stale := store.OtherVectorSlot(g.Slot()) + if empty, err := ps.VectorSlotEmpty(ctx, stale); err == nil && !empty { + t0 := time.Now() + if err := ps.ClearVectorSlot(ctx, stale); err != nil { + log.Printf("pebble-serve: clearing stale HNSW slot %#x failed: %v", stale, err) + } else { + log.Printf("pebble-serve: cleared stale HNSW slot %#x in %s", stale, time.Since(t0).Round(time.Millisecond)) + } + } } // startInProcessCrawl wires the crawler-inside-serve flow. Bumps @@ -973,78 +985,76 @@ func (s *pebbleHTTP) startInProcessCrawl(ctx context.Context, ps *store.PebbleSt } _ = c.Seed(u) } - log.Printf("in-serve crawler: %d seeds queued (concurrency=%d, depth=%d, checkpoint=%s)", - len(seeds), cfg.Crawler.MaxConcurrent, cfg.Crawler.MaxDepth, ckpEvery) + log.Printf("in-serve crawler: %d seeds queued (concurrency=%d, depth=%d, checkpoint=%s, zombie-reclaim=%v)", + len(seeds), cfg.Crawler.MaxConcurrent, cfg.Crawler.MaxDepth, ckpEvery, crawler.ZombieReclaimEnabled()) s.crawlActive = true - // Checkpoint goroutine: incremental persist via PersistFrom — each - // tick only writes nodes [lastN, n). Shutdown does a full persist - // from 0 so any backlinks added to older nodes get refreshed. - // Seeding lastN from the loaded graph means the first checkpoint - // after restart only writes nodes added during this run; the prior - // N are already on disk. + // Checkpoint goroutine: incremental persist via TryPersistFrom — each + // tick writes nodes [lastN, n) plus the dirty set (invalidated nodes and + // nodes that gained back-links), so shutdown is just one more incremental + // checkpoint. Seeding lastN from the loaded graph means the first + // checkpoint after restart only writes nodes added during this run. wg.Add(1) go func() { defer wg.Done() t := time.NewTicker(ckpEvery) defer t.Stop() lastN := g.Len() + pqLastN := lastN if lastN > 0 { log.Printf("in-serve crawler: checkpoint baseline = %d nodes (loaded from disk)", lastN) } + checkpoint := func(what string) bool { + n, dirty := g.Len(), g.DirtyCount() + if n == 0 || (n == lastN && dirty == 0) { + return false + } + // graph can shrink (e.g., /admin/hnsw-compact rewrites + // indices and writes a smaller meta). When that happens, lastN + // from before the compaction is stale and > n; PersistFrom(lastN) + // would be a no-op forever, stranding any new AddPassages until + // shutdown. The compact job does its own full persist so disk + // is already in sync; we just need to resync lastN here. + if n < lastN { + lastN = n + return false + } + t0 := time.Now() + ok, err := g.TryPersistFrom(context.Background(), ps, lastN) + if err != nil { + log.Printf("in-serve crawler: HNSW %s (incremental from %d) failed: %v", what, lastN, err) + return false + } + if !ok { + log.Printf("in-serve crawler: HNSW %s skipped — full persist in progress", what) + return false + } + log.Printf("in-serve crawler: HNSW %s at %d nodes (+%d new, +%d dirty, took %s)", + what, n, n-lastN, dirty, time.Since(t0)) + lastN = n + return true + } for { select { case <-ctx.Done(): - n := g.Len() - if n > 0 { - t0 := time.Now() - log.Printf("in-serve crawler: final HNSW persist at shutdown (%d nodes, full)", n) - if err := g.Persist(context.Background(), ps); err != nil { - log.Printf("in-serve crawler: final HNSW persist failed: %v", err) - } else { - log.Printf("in-serve crawler: final HNSW persist complete in %s", time.Since(t0)) - } - } + checkpoint("final checkpoint at shutdown") return case <-t.C: - n := g.Len() - if n == 0 || n == lastN { - continue - } - // graph can shrink (e.g., /admin/hnsw-compact rewrites - // indices and writes a smaller meta). When that happens, lastN - // from before the compaction is stale and > n; PersistFrom(lastN) - // would be a no-op forever, stranding any new AddPassages until - // shutdown. The compact handler does its own full Persist so disk - // is already in sync; we just need to resync lastN here. - if n < lastN { - lastN = n - continue - } - t0 := time.Now() - if err := g.PersistFrom(context.Background(), ps, lastN); err != nil { - log.Printf("in-serve crawler: HNSW persist (incremental from %d) failed: %v", lastN, err) + if !checkpoint("checkpoint") { continue } + n := g.Len() // alongside HNSW node writes, persist any newly- // encoded PQ codes for nodes [lastN, n). Skipped silently // when no codebook is loaded. - pqWritten := 0 if g.HasPQ() { - if w, err := g.PersistPQCodesFrom(context.Background(), ps, lastN); err != nil { + if w, err := g.PersistPQCodesFrom(context.Background(), ps, pqLastN); err != nil { log.Printf("in-serve crawler: PQ codes persist failed: %v", err) - } else { - pqWritten = w + } else if w > 0 { + log.Printf("in-serve crawler: +%d PQ codes persisted", w) } } - if pqWritten > 0 { - log.Printf("in-serve crawler: HNSW checkpoint at %d nodes (+%d incremental, +%d PQ codes, took %s)", - n, n-lastN, pqWritten, time.Since(t0)) - } else { - log.Printf("in-serve crawler: HNSW checkpoint at %d nodes (+%d incremental, took %s)", - n, n-lastN, time.Since(t0)) - } - lastN = n + pqLastN = n } } }() @@ -1263,6 +1273,10 @@ type pebbleHTTP struct { hnswTotal atomic.Uint64 // Nodes invalidated by the post-load store reconcile (purge orphans). reconciledOrphans atomic.Int64 + // Background admin jobs (hnsw-compact) that touch the store; joined + // before the store closes. + bgJobs sync.WaitGroup + compact compactJob // Dense/hybrid candidates that failed GetDocByURL resolution and were // silently dropped from responses — the store/graph divergence signal. denseResolutionDrops atomic.Int64 diff --git a/cmd/cosift/serve_stats.go b/cmd/cosift/serve_stats.go index 9966395..c4cc9be 100644 --- a/cmd/cosift/serve_stats.go +++ b/cmd/cosift/serve_stats.go @@ -397,6 +397,11 @@ func (s *pebbleHTTP) buildStatsBody(ctx context.Context) ([]byte, error) { out["hnsw_loaded"] = s.hnsw() != nil // async load progress (state/pct/ETA) so a restart's warm-up is watchable. out["hnsw_load"] = s.hnswLoadSnapshot() + out["hnsw_compact"] = s.compact.snapshot() + if g := s.hnsw(); g != nil { + out["hnsw_reclaimed_total"] = g.Reclaimed() + out["hnsw_slot"] = g.Slot() + } // PQ status — operator-facing visibility into compression // state. Only present when the graph is loaded; nil otherwise. if s.hnsw() != nil { @@ -582,6 +587,22 @@ func (s *pebbleHTTP) handleMetrics(w http.ResponseWriter, r *http.Request) { fmt.Fprintf(w, "# TYPE cosift_vector_dim gauge\n") fmt.Fprintf(w, "cosift_vector_dim %d\n", s.vectorDim) } + if g := s.hnsw(); g != nil { + st := g.PQStatus() + fmt.Fprintf(w, "# HELP cosift_hnsw_zombie_nodes HNSW nodes invalidated (reclaimed or reconciled) and awaiting compaction.\n") + fmt.Fprintf(w, "# TYPE cosift_hnsw_zombie_nodes gauge\n") + fmt.Fprintf(w, "cosift_hnsw_zombie_nodes %d\n", st.NodesTotal-st.NodesValid) + fmt.Fprintf(w, "# HELP cosift_hnsw_reclaimed_total Nodes invalidated by re-crawl zombie reclaim since process start.\n") + fmt.Fprintf(w, "# TYPE cosift_hnsw_reclaimed_total counter\n") + fmt.Fprintf(w, "cosift_hnsw_reclaimed_total %d\n", g.Reclaimed()) + running := 0 + if s.compact.snapshot()["state"] == "running" { + running = 1 + } + fmt.Fprintf(w, "# HELP cosift_hnsw_compact_running 1 while an hnsw-compact job is in flight.\n") + fmt.Fprintf(w, "# TYPE cosift_hnsw_compact_running gauge\n") + fmt.Fprintf(w, "cosift_hnsw_compact_running %d\n", running) + } // PromQL // rate(cosift_request_duration_seconds_sum) / rate(cosift_requests_total) // gives mean latency in any window. Labels = path; misrouted calls (404) diff --git a/internal/crawler/crawler.go b/internal/crawler/crawler.go index 57fe2f2..caf7fe5 100644 --- a/internal/crawler/crawler.go +++ b/internal/crawler/crawler.go @@ -763,7 +763,7 @@ func (c *Crawler) runEmbedJob(parent context.Context, job *embedJob, jobTimeout } // Mirror the synchronous path's zombie reclaim so re-crawled // URLs don't accumulate generations of vectors in HNSW. - if os.Getenv("COSIFT_ZOMBIE_RECLAIM") == "1" { + if ZombieReclaimEnabled() { if inv, ok := c.passageWriter.(URLInvalidator); ok { _, _ = inv.MarkURLInvalid(ctx, job.url) } @@ -1480,14 +1480,12 @@ func (c *Crawler) processClaimed(ctx context.Context, item store.FrontierItem, g // writes give readers more chances to slip in. Same total // lock time, smaller bursts. if c.passageWriter != nil { - // When this URL was - // previously crawled, the prior generation of chunks - // still lives in the HNSW graph (same url, stale vecs). - // Mark them invalid before adding the fresh set so the - // graph doesn't accumulate generations. Gated by env - // COSIFT_ZOMBIE_RECLAIM=1 until soaked; off-by-default - // preserves prior behavior bit-for-bit. - if os.Getenv("COSIFT_ZOMBIE_RECLAIM") == "1" { + // When this URL was previously crawled, the prior + // generation of chunks still lives in the HNSW graph + // (same url, stale vecs). Mark them invalid before + // adding the fresh set so the graph doesn't accumulate + // generations. + if ZombieReclaimEnabled() { inv, ok := c.passageWriter.(URLInvalidator) if !ok { log.Printf("zombie-reclaim: passageWriter %T does NOT implement URLInvalidator (one-time check)", c.passageWriter) diff --git a/internal/crawler/store_iface.go b/internal/crawler/store_iface.go index d73189e..544fd19 100644 --- a/internal/crawler/store_iface.go +++ b/internal/crawler/store_iface.go @@ -9,6 +9,8 @@ package crawler import ( "context" + "os" + "strings" "github.com/pilot-protocol/cosift/internal/store" ) @@ -104,6 +106,16 @@ type PassageWriterBatch interface { UpsertPassageBatch(ctx context.Context, ps []*store.Passage) error } +// ZombieReclaimEnabled gates re-crawl reclaim of prior HNSW generations: +// on unless COSIFT_ZOMBIE_RECLAIM is "0", "false" or "off". +func ZombieReclaimEnabled() bool { + switch strings.ToLower(os.Getenv("COSIFT_ZOMBIE_RECLAIM")) { + case "0", "false", "off": + return false + } + return true +} + // URLInvalidator is the optional zombie-reclaim surface. When a doc is // re-crawled (URL already in famDoc), the old passage vectors persist in // the HNSW graph alongside the new ones — same URL, multiple generations diff --git a/internal/crawler/wet.go b/internal/crawler/wet.go index c6537e1..cb9d9f4 100644 --- a/internal/crawler/wet.go +++ b/internal/crawler/wet.go @@ -305,10 +305,8 @@ func (c *Crawler) indexWetRecord(ctx context.Context, rec *WetRecord, lexicalOnl if embErr != nil || len(vecs) != len(chunks) { return nil // best-effort — BM25 doc is already indexed } - if inv, ok := c.passageWriter.(URLInvalidator); ok { - if getEnv("COSIFT_ZOMBIE_RECLAIM") == "1" { - _, _ = inv.MarkURLInvalid(ctx, rec.URL) - } + if inv, ok := c.passageWriter.(URLInvalidator); ok && ZombieReclaimEnabled() { + _, _ = inv.MarkURLInvalid(ctx, rec.URL) } for i, ch := range chunks { p := &store.Passage{ diff --git a/internal/crawler/zombie_reclaim_test.go b/internal/crawler/zombie_reclaim_test.go new file mode 100644 index 0000000..2090e53 --- /dev/null +++ b/internal/crawler/zombie_reclaim_test.go @@ -0,0 +1,15 @@ +package crawler + +import "testing" + +func TestZombieReclaimEnabled(t *testing.T) { + for _, tc := range []struct { + v string + want bool + }{{"", true}, {"1", true}, {"true", true}, {"0", false}, {"false", false}, {"OFF", false}} { + t.Setenv("COSIFT_ZOMBIE_RECLAIM", tc.v) + if got := ZombieReclaimEnabled(); got != tc.want { + t.Errorf("%q: got %v want %v", tc.v, got, tc.want) + } + } +} diff --git a/internal/index/hnsw_compact.go b/internal/index/hnsw_compact.go index 83eed7f..e1fd1e6 100644 --- a/internal/index/hnsw_compact.go +++ b/internal/index/hnsw_compact.go @@ -1,6 +1,13 @@ package index -import "log" +import ( + "context" + "fmt" + "log" + "time" + + "github.com/pilot-protocol/cosift/internal/store" +) // compactProgressEvery paces the in-compact progress logs; the whole pass // runs under the write lock, so these lines are the only liveness signal. @@ -61,6 +68,92 @@ func (h *HNSW) Rebuild() *HNSW { // even without PQ. Compaction restores the recall the underlying corpus // can support. func (h *HNSW) Compact() (removed int) { + h.persistMu.Lock() + defer h.persistMu.Unlock() + return h.compactLocked() +} + +// DirtyCount reports how many nodes await an incremental persist. +func (h *HNSW) DirtyCount() int { + h.mu.RLock() + defer h.mu.RUnlock() + return len(h.dirty) +} + +// CompactProgress is the observable state of a CompactPersist run. +type CompactProgress struct { + Phase string // compact | persist | cleanup | done | error + NodesBefore, NodesAfter, Removed int + Written, Total int // persist progress +} + +// CompactResult summarizes a finished CompactPersist. +type CompactResult struct { + NodesBefore, NodesAfter, Removed int + CompactDur, PersistDur time.Duration + Persisted, Forced bool +} + +// CompactPersist compacts the graph and, unless skipPersist (or nothing was +// removed and !forcePersist), rewrites it into the inactive slot, then clears +// the previous slot and the PQ family (codes are keyed by node id). No other +// persist runs in between. progress may be nil. +func (h *HNSW) CompactPersist(ctx context.Context, ps *store.PebbleStore, skipPersist, forcePersist bool, progress func(CompactProgress)) (CompactResult, error) { + h.persistMu.Lock() + defer h.persistMu.Unlock() + report := func(p CompactProgress) { + if progress != nil { + progress(p) + } + } + res := CompactResult{NodesBefore: h.Len()} + report(CompactProgress{Phase: "compact", NodesBefore: res.NodesBefore}) + t0 := time.Now() + res.Removed = h.compactLocked() + res.CompactDur = time.Since(t0) + res.NodesAfter = h.Len() + base := CompactProgress{NodesBefore: res.NodesBefore, NodesAfter: res.NodesAfter, Removed: res.Removed} + if skipPersist || (res.Removed == 0 && !forcePersist) { + base.Phase = "done" + report(base) + return res, nil + } + res.Forced = forcePersist && res.Removed == 0 + + base.Phase = "persist" + report(base) + t0 = time.Now() + old := h.Slot() + if err := h.persistSwapLocked(ctx, ps, func(p PersistProgress) { + pp := base + pp.Written, pp.Total = p.Written, p.Total + report(pp) + }); err != nil { + base.Phase = "error" + report(base) + return res, err + } + res.PersistDur = time.Since(t0) + res.Persisted = true + + base.Phase = "cleanup" + report(base) + if err := ps.ClearVectorSlot(ctx, old); err != nil { + base.Phase = "error" + report(base) + return res, fmt.Errorf("clear old slot: %w", err) + } + if err := ps.ClearPQFamily(ctx); err != nil { + base.Phase = "error" + report(base) + return res, fmt.Errorf("clear pq family: %w", err) + } + base.Phase = "done" + report(base) + return res, nil +} + +func (h *HNSW) compactLocked() (removed int) { h.mu.Lock() defer h.mu.Unlock() diff --git a/internal/index/hnsw_persist.go b/internal/index/hnsw_persist.go index 6d23665..d4369f8 100644 --- a/internal/index/hnsw_persist.go +++ b/internal/index/hnsw_persist.go @@ -108,6 +108,10 @@ func (h *HNSW) TryPersistFrom(ctx context.Context, ps *store.PebbleStore, fromId func (h *HNSW) PersistSwap(ctx context.Context, ps *store.PebbleStore, progress func(PersistProgress)) error { h.persistMu.Lock() defer h.persistMu.Unlock() + return h.persistSwapLocked(ctx, ps, progress) +} + +func (h *HNSW) persistSwapLocked(ctx context.Context, ps *store.PebbleStore, progress func(PersistProgress)) error { h.mu.RLock() old := h.slot h.mu.RUnlock() From 486f8e45b1b422202b7df62bd7d510f8d5adee47 Mon Sep 17 00:00:00 2001 From: Andrei Marinescu Date: Wed, 2 Sep 2026 11:01:30 +0300 Subject: [PATCH 3/3] deploy/docs: compact timer + threshold 15 with async-job polling; ENV/PEBBLE docs for reclaim default, slots, compaction knob --- deploy/README.md | 1 + deploy/harvesters/cosift-compact.sh | 64 ++++++++++++++++++++------- deploy/scripts/cosift-self-update.sh | 5 ++- deploy/systemd/cosift-compact.service | 10 +++++ deploy/systemd/cosift-compact.timer | 8 ++++ docs/ENV.md | 3 +- docs/PEBBLE.md | 10 ++++- 7 files changed, 79 insertions(+), 22 deletions(-) create mode 100644 deploy/systemd/cosift-compact.service create mode 100644 deploy/systemd/cosift-compact.timer diff --git a/deploy/README.md b/deploy/README.md index 36b1916..0a6d64d 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -24,6 +24,7 @@ content pipeline is reproducible from git, not only on the box. | wet-refresh | CommonCrawl WET (rotating offset, lexical-only) | daily 03:00 | | pypi/hf-task/github-topic | popular PyPI / HF-by-task / GitHub-by-topic | weekly | | sitemap-refresh | (disabled — net-negative on single box) | — | +| cosift-compact | `POST /admin/hnsw-compact` when `/stats` zombies ≥ 15% (async job, polls `hnsw_compact`) | weekly Sun 04:30 | ## Notes - Curated ingester: `disable_link_following=true` + `include_domains` allowlist in diff --git a/deploy/harvesters/cosift-compact.sh b/deploy/harvesters/cosift-compact.sh index b87a569..7ee55a2 100755 --- a/deploy/harvesters/cosift-compact.sh +++ b/deploy/harvesters/cosift-compact.sh @@ -1,29 +1,59 @@ #!/bin/bash -# cosift-compact.sh — invoke hnsw-compact when zombies exceed a threshold. -# Run from cosift-compact.timer (weekly recommended). -set -e -TS=$(date -u +%Y-%m-%dT%H:%M:%SZ) -THRESHOLD_PCT=30 -# /admin/hnsw-compact is gated by cluster.peer_auth_token — read it from the -# same config the server uses so this keeps working once the token is set. +# cosift-compact.sh — start hnsw-compact when zombies exceed a threshold and +# follow the async job through /stats.hnsw_compact. Run from cosift-compact.timer. +set -u +TS() { date -u +%Y-%m-%dT%H:%M:%SZ; } +LOG=/var/log/cosift-compact.log +log() { echo "$(TS) $*" | sudo tee -a "$LOG" >/dev/null; } +THRESHOLD_PCT=${THRESHOLD_PCT:-15} +POLL_SEC=${POLL_SEC:-30} +MAX_WAIT_SEC=${MAX_WAIT_SEC:-28800} +STATS_URL=http://127.0.0.1:7777/stats TOKEN=$(python3 -c "import json; print(json.load(open('/home/ubuntu/cosift.json')).get('cluster',{}).get('peer_auth_token',''))" 2>/dev/null || echo "") -STATS=$(curl -s --max-time 30 http://127.0.0.1:7777/stats) + +STATS=$(curl -s --max-time 30 "$STATS_URL") if [ -z "$STATS" ]; then - echo "${TS} stats fetch failed" | sudo tee -a /var/log/cosift-compact.log >/dev/null + log "stats fetch failed" exit 0 fi PCT=$(echo "$STATS" | python3 -c " import json,sys d=json.load(sys.stdin) -t=d[\"pq\"][\"nodes_total\"] -z=d[\"pq\"][\"zombie_nodes\"] -if t==0: print(0); exit() -print(int(100*z/t)) +pq=d.get('pq') +if not pq: print(-1); exit() +t=pq['nodes_total']; z=pq['zombie_nodes'] +print(int(100*z/t) if t else 0) ") -echo "${TS} zombies=${PCT}% threshold=${THRESHOLD_PCT}%" | sudo tee -a /var/log/cosift-compact.log >/dev/null +if [ "$PCT" -lt 0 ]; then + log "graph not loaded (no pq in /stats), skipping" + exit 0 +fi +log "zombies=${PCT}% threshold=${THRESHOLD_PCT}%" if [ "$PCT" -lt "$THRESHOLD_PCT" ]; then - echo "${TS} below threshold, skipping compact" | sudo tee -a /var/log/cosift-compact.log >/dev/null + log "below threshold, skipping compact" exit 0 fi -RESULT=$(curl -s --max-time 1800 -X POST ${TOKEN:+-H "Authorization: Bearer $TOKEN"} http://127.0.0.1:7777/admin/hnsw-compact) -echo "${TS} compact result: ${RESULT}" | sudo tee -a /var/log/cosift-compact.log >/dev/null + +START=$(curl -s --max-time 30 -o /dev/stderr -w '%{http_code}' -X POST ${TOKEN:+-H "Authorization: Bearer $TOKEN"} http://127.0.0.1:7777/admin/hnsw-compact 2>&1) +CODE=${START: -3} +if [ "$CODE" != "202" ]; then + log "compact start returned HTTP ${CODE}: ${START%???}" + exit 0 +fi +log "compact started" +DEADLINE=$(( $(date +%s) + MAX_WAIT_SEC )) +while [ "$(date +%s)" -lt "$DEADLINE" ]; do + sleep "$POLL_SEC" + JOB=$(curl -s --max-time 30 "$STATS_URL" | python3 -c " +import json,sys +j=json.load(sys.stdin).get('hnsw_compact',{}) +print(j.get('state','?'), json.dumps(j, sort_keys=True)) +" 2>/dev/null) + STATE=${JOB%% *} + case "$STATE" in + running) log "running: ${JOB#* }" ;; + done|error) log "compact ${STATE}: ${JOB#* }"; exit 0 ;; + *) log "stats unavailable (${STATE}), still waiting" ;; + esac +done +log "gave up waiting after ${MAX_WAIT_SEC}s; job may still be running" diff --git a/deploy/scripts/cosift-self-update.sh b/deploy/scripts/cosift-self-update.sh index 43d81c1..d00005e 100755 --- a/deploy/scripts/cosift-self-update.sh +++ b/deploy/scripts/cosift-self-update.sh @@ -180,8 +180,9 @@ log "restarting $SERVICE" sudo systemctl restart "$SERVICE" # --- HEALTH GATE --------------------------------------------------------- -# The listener only binds AFTER the ~4-5 min synchronous HNSW load, so a -# 200 from /healthz is a true readiness signal. Poll up to the timeout. +# The listener binds immediately (BM25 serves while the HNSW graph loads in +# the background), so a 200 from /healthz means the new binary is up, not +# that dense retrieval is ready — watch /stats.hnsw_load for that. log "health-gating $HEALTH_URL (timeout ${HEALTH_TIMEOUT_S}s, every ${HEALTH_INTERVAL_S}s)" deadline=$(( $(date +%s) + HEALTH_TIMEOUT_S )) healthy=0 diff --git a/deploy/systemd/cosift-compact.service b/deploy/systemd/cosift-compact.service new file mode 100644 index 0000000..0f9c52c --- /dev/null +++ b/deploy/systemd/cosift-compact.service @@ -0,0 +1,10 @@ +[Unit] +Description=Cosift HNSW compact (zombie threshold check + async job follow) +After=cosift-serve.service +Wants=cosift-serve.service +[Service] +Type=oneshot +User=ubuntu +ExecStart=/usr/local/bin/cosift-compact.sh +StandardOutput=journal +StandardError=journal diff --git a/deploy/systemd/cosift-compact.timer b/deploy/systemd/cosift-compact.timer new file mode 100644 index 0000000..b18f868 --- /dev/null +++ b/deploy/systemd/cosift-compact.timer @@ -0,0 +1,8 @@ +[Unit] +Description=Cosift HNSW compact schedule +[Timer] +OnCalendar=Sun *-*-* 04:30:00 +Persistent=true +AccuracySec=10min +[Install] +WantedBy=timers.target diff --git a/docs/ENV.md b/docs/ENV.md index 08f824c..9c610ca 100644 --- a/docs/ENV.md +++ b/docs/ENV.md @@ -200,7 +200,7 @@ not re-enable them without the confidentiality decision in that section. | `COSIFT_CRAWL_PDF` | bool (`"false"` disables) | unset → PDF parsing **enabled** (sandboxed) | Set to `"false"` to disable sandboxed PDF parsing. Any other value leaves it on. | `internal/crawler/crawler.go:1155` | | `COSIFT_REFETCH_AFTER_HOURS` | int (hours) | `0` → disabled (every revisit issues a conditional GET) | Skip re-fetching a healthy URL fetched within this window. Also defines the "fresh" window for prefer-new (defaults to 24h there) and the WET fresh window. Must be `> 0`. | `crawler.go:1107,1492`; `wet.go:94` | | `COSIFT_PREFER_NEW_URLS` | bool (`"1"`) | unset → off | When `"1"`, the frontier prefers never-seen URLs over recently-fetched ones (one `GetDocByURL` per candidate). | `internal/crawler/crawler.go:1488` | -| `COSIFT_ZOMBIE_RECLAIM` | bool (`"1"`) | unset → off | When `"1"`, marks a re-crawled URL's prior chunk generation invalid in the HNSW graph before adding fresh vectors. **Do not enable on large graphs**: each reclaim is a full O(N) scan under the exclusive graph write lock (~1 s at 80M nodes, per re-crawled doc, from concurrent embed workers) — search and ingest stall. Needs a URL→node index first. | `crawler.go:766,1490`; `wet.go:308` | +| `COSIFT_ZOMBIE_RECLAIM` | bool (`"0"`/`"false"`/`"off"` disables) | unset → **on** | Marks a re-crawled URL's prior chunk generation invalid in the HNSW graph before adding fresh vectors, so the graph doesn't accumulate generations. O(k) per URL via the graph's URL index; invalidations persist with the next checkpoint. Zombies accumulate at the re-crawl rate until `/admin/hnsw-compact` (weekly timer, 15% threshold). Counters: `/stats.hnsw_reclaimed_total`, `cosift_hnsw_zombie_nodes`. | `crawler.go` (`ZombieReclaimEnabled`) | | `COSIFT_EMBED_DECOUPLE_WORKERS` | int | `0` → decoupled embed pipeline **off** | Number of dedicated embed-worker goroutines that drain the crawl embed queue (decouples crawl from embed/HNSW-write latency). `0` keeps the synchronous path. Must be `>= 0`. | `internal/crawler/crawler.go:586` | | `COSIFT_EMBED_DECOUPLE_BUFFER` | int | `4096` | Buffer size of the decoupled embed queue (only when workers `> 0`). Must be `>= 0`. | `internal/crawler/crawler.go:588` | | `COSIFT_HOSTSWEEP_DISABLED` | bool (`"1"`) | unset → sweeper **enabled** | When `"1"`, disables the self-cleaning host sweeper entirely. | `internal/crawler/crawler.go:868` | @@ -225,6 +225,7 @@ not re-enable them without the confidentiality decision in that section. | `COSIFT_PEBBLE_CACHE_MB` | int (MB) | `128` | Pebble block-cache size in MB; must be `> 0`. | `internal/store/pebble.go:153` | | `COSIFT_PEBBLE_MEMTABLE_MB` | int (MB) | `32` | Pebble memtable size in MB; must be `> 0`. | `internal/store/pebble.go:154` | | `COSIFT_PEBBLE_MEMTABLES` | int | `2` | Memtable count; `MemTableStopWritesThreshold = value + 2`. Must be `> 0`. | `internal/store/pebble.go:155` | +| `COSIFT_PEBBLE_COMPACTIONS` | int | `1` | Pebble `MaxConcurrentCompactions`. `1` (Pebble's default) serializes every background compaction on one slot — the cause of the 128K-SSTable pile and the 12 MB/s full persist observed on the production box; `4` is a sane value on a 64-core host. Must be `> 0`. | `internal/store/pebble.go` (`openPebble`) | | `COSIFT_PEBBLE_SYNC` | bool (`"false"` disables) | unset → `Sync` (fsync each commit) | Set to `"false"` to use `NoSync` writes (skips per-commit fsync — faster crawls, drops durability vs OS crash; WAL still written so process-crash durability holds). | `internal/store/pebble.go:176` | --- diff --git a/docs/PEBBLE.md b/docs/PEBBLE.md index c6333b6..d6d571e 100644 --- a/docs/PEBBLE.md +++ b/docs/PEBBLE.md @@ -28,8 +28,14 @@ The two backends have parity for BM25 search quality (`TestPebbleBM25MatchesSQLi 'f' + 'u' + url → packed frontierEntry 'f' + 'q' + host + 0x00 + url → empty (queued, host-keyed index) 'f' + 'i' + host + 0x00 + url → empty (in-flight, host-keyed index) -'v' + 0x00 + "meta" → HNSW meta blob -'v' + 0x01 + uint64-be(nodeID) → HNSW node blob +'v' + 0x00 + "meta" → HNSW meta blob (HSW1 = 20 B, slot 0x01 implied; HSW2 = 21 B, trailing slot byte) +'v' + 0x01 + uint64-be(nodeID) → HNSW node blob, slot A +'v' + 0x02 + uint64-be(nodeID) → HNSW node blob, slot B + +A full persist (`hnsw-compact`, `hnsw-rebuild`) writes the next generation into +the inactive slot, then points meta at it, then clears the old slot — the previous +graph stays loadable throughout and the fresh writes never land on tombstones. +Incremental checkpoints write new + dirtied nodes into the active slot. 'm' + name → counter bytes (next_doc_id, next_term_id, sum_doc_len, indexed_docs) ```