diff --git a/internal/index/hnsw_compact.go b/internal/index/hnsw_compact.go index ba4b09f..9fcf13e 100644 --- a/internal/index/hnsw_compact.go +++ b/internal/index/hnsw_compact.go @@ -1,5 +1,12 @@ package index +import "log" + +// compactProgressEvery paces the in-compact progress logs; the whole pass +// runs under the write lock, so these lines are the only liveness signal. +// A var (not const) so tests can lower it below the fixture size. +var compactProgressEvery = 10_000_000 + // Rebuild constructs a fresh HNSW with the same parameters as h and inserts // every valid (vec != nil) node from h via AddPassage. Unlike Compact, which // merely removes zombies and rewires existing edges, Rebuild reconstructs @@ -69,6 +76,9 @@ func (h *HNSW) Compact() (removed int) { newCodes = make([][]uint16, 0, len(h.codes)) } for i := range h.nodes { + if i > 0 && i%compactProgressEvery == 0 { + log.Printf("hnsw compact: scanning %d/%d nodes", i, len(h.nodes)) + } if len(h.nodes[i].vec) == 0 { remap[i] = -1 continue @@ -88,6 +98,9 @@ func (h *HNSW) Compact() (removed int) { // 2. Remap neighbor lists. Iterating in increasing order so writes only // touch slots we've already read from the source array. for i := range newNodes { + if i > 0 && i%compactProgressEvery == 0 { + log.Printf("hnsw compact: rewiring neighbors %d/%d nodes", i, len(newNodes)) + } for lvl := range newNodes[i].neighbors { old := newNodes[i].neighbors[lvl] out := make([]int, 0, len(old)) diff --git a/internal/index/hnsw_compact_test.go b/internal/index/hnsw_compact_test.go index 1eea15f..ece45cd 100644 --- a/internal/index/hnsw_compact_test.go +++ b/internal/index/hnsw_compact_test.go @@ -1,9 +1,13 @@ package index import ( + "bytes" "context" "fmt" + "log" "math/rand" + "os" + "strings" "testing" ) @@ -118,3 +122,32 @@ func TestHNSWZombieCompaction(t *testing.T) { t.Errorf("rebuild did not restore recall: clean=%.3f rebuilt=%.3f", clean, rebuiltRecall) } } + +// TestHNSWCompactProgressLogs pins the write-lock liveness lines. +func TestHNSWCompactProgressLogs(t *testing.T) { + h := buildTestHNSW(250, 16, 3, 5) + for i := 0; i < 40; i++ { + h.nodes[i*5].vec = nil + } + + savedEvery := compactProgressEvery + compactProgressEvery = 100 + var buf bytes.Buffer + log.SetOutput(&buf) + defer func() { + compactProgressEvery = savedEvery + log.SetOutput(os.Stderr) + }() + + removed := h.Compact() + if removed != 40 { + t.Errorf("removed: want 40, got %d", removed) + } + out := buf.String() + if !strings.Contains(out, "hnsw compact: scanning") { + t.Errorf("missing scanning progress line in:\n%s", out) + } + if !strings.Contains(out, "hnsw compact: rewiring neighbors") { + t.Errorf("missing rewiring progress line in:\n%s", out) + } +} diff --git a/internal/index/hnsw_persist.go b/internal/index/hnsw_persist.go index f89f5f9..837dd33 100644 --- a/internal/index/hnsw_persist.go +++ b/internal/index/hnsw_persist.go @@ -41,12 +41,21 @@ import ( "fmt" "log" "math" + "time" "github.com/pilot-protocol/cosift/internal/store" ) const hnswMetaMagic = "HSW1" +// persistWindowBytes bounds encoded blobs held in memory at once: a full +// persist that materializes every blob first costs ~vec-bytes of extra heap +// (~240 GB at 80M nodes) and OOMs before writing anything. Var for tests. +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 @@ -55,7 +64,7 @@ func (h *HNSW) Persist(ctx context.Context, ps *store.PebbleStore) error { return h.PersistFrom(ctx, ps, 0) } -// PersistFrom writes meta + nodes[fromIdx:] in a single Pebble batch. The +// 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. @@ -81,15 +90,48 @@ func (h *HNSW) PersistFrom(ctx context.Context, ps *store.PebbleStore, fromIdx i // 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). - entries := make([]store.VectorNodeEntry, 0, len(h.nodes)-fromIdx) + total := len(h.nodes) - fromIdx + start := time.Now() + window := make([]store.VectorNodeEntry, 0, 4096) + windowBytes, written, flushes := 0, 0, 0 + var bytesWritten int64 + flush := func() error { + if len(window) == 0 { + return nil + } + if err := ps.PutVectorNodesBatch(ctx, window); err != nil { + return fmt.Errorf("put vector nodes batch: %w", err) + } + written += len(window) + bytesWritten += int64(windowBytes) + flushes++ + if persistFlushed != nil { + persistFlushed(len(window), windowBytes) + } + if flushes > 1 || written < total { + elapsed := max(time.Since(start).Seconds(), 0.001) + rate := float64(written) / elapsed + eta := time.Duration(float64(total-written) / rate * float64(time.Second)).Round(time.Second) + log.Printf("hnsw persist: %d/%d nodes (%.1f GiB, %.0f nodes/s, eta %s)", + written, total, float64(bytesWritten)/(1<<30), rate, eta) + } + clear(window) + window = window[:0] + windowBytes = 0 + return nil + } for i := fromIdx; i < len(h.nodes); i++ { - entries = append(entries, store.VectorNodeEntry{ - ID: uint64(i), - Blob: encodeHNSWNode(&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 + } + } } - if err := ps.PutVectorNodesBatch(ctx, entries); err != nil { - return fmt.Errorf("put vector nodes batch: %w", err) + 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 { diff --git a/internal/index/hnsw_persist_test.go b/internal/index/hnsw_persist_test.go index 937318e..ab16606 100644 --- a/internal/index/hnsw_persist_test.go +++ b/internal/index/hnsw_persist_test.go @@ -6,6 +6,7 @@ import ( "math/rand" "path/filepath" "sort" + "strings" "testing" "github.com/pilot-protocol/cosift/internal/store" @@ -496,3 +497,252 @@ func sameIntSlice(a, b []int) bool { } return true } + +// failAfterFlagCtx reports Canceled once *fail is set; used to kill a persist +// between windows deterministically. +type failAfterFlagCtx struct { + context.Context + fail *bool +} + +func (c *failAfterFlagCtx) Err() error { + if *c.fail { + return context.Canceled + } + return nil +} + +func buildTestHNSW(n, dim int, graphSeed, vecSeed int64) *HNSW { + h := NewHNSW(dim) + h.rng = rand.New(rand.NewSource(graphSeed)) + rng := rand.New(rand.NewSource(vecSeed)) + for i := 0; i < n; i++ { + v := make([]float32, dim) + for j := range v { + v[j] = float32(rng.NormFloat64()) + } + h.AddPassage(fmt.Sprintf("https://x/%d", i), fmt.Sprintf("doc %d", i), i*100, 50, v) + } + return h +} + +// TestHNSWPersistWindowed pins the streaming persist: a multi-window persist +// must keep every window bounded, cover every node exactly once, and load +// back identical to a single-window persist. +func TestHNSWPersistWindowed(t *testing.T) { + dir := filepath.Join(t.TempDir(), "pebble") + ps, err := store.OpenPebble(dir) + if err != nil { + t.Fatalf("OpenPebble: %v", err) + } + defer ps.Close() + ctx := context.Background() + + const ( + n = 500 + dim = 32 + ) + original := buildTestHNSW(n, dim, 13, 7) + + savedWindow := persistWindowBytes + persistWindowBytes = 16 * 1024 + windows := 0 + nodesFlushed := 0 + persistFlushed = func(nodes, bytes int) { + windows++ + nodesFlushed += nodes + // Window may overshoot by at most the final blob (+16 overhead). + if bytes > persistWindowBytes+8*1024 { + t.Errorf("window %d: %d bytes exceeds bound", windows, bytes) + } + } + defer func() { + persistWindowBytes = savedWindow + persistFlushed = nil + }() + + if err := original.Persist(ctx, ps); err != nil { + t.Fatalf("persist: %v", err) + } + if windows < 3 { + t.Fatalf("windows: want >=3 for this fixture, got %d", windows) + } + if nodesFlushed != n { + t.Fatalf("nodes flushed: want %d, got %d", n, nodesFlushed) + } + + loaded, ok, err := LoadHNSW(ctx, ps) + if err != nil { + t.Fatalf("load: %v", err) + } + if !ok { + t.Fatal("loaded ok=false; expected persisted index") + } + if loaded.Len() != n { + t.Errorf("loaded.Len(): want %d, got %d", n, loaded.Len()) + } + if loaded.entryPoint != original.entryPoint || loaded.maxLevel != original.maxLevel { + t.Errorf("meta mismatch: entry %d/%d level %d/%d", + loaded.entryPoint, original.entryPoint, loaded.maxLevel, original.maxLevel) + } + rng := rand.New(rand.NewSource(99)) + for qi := 0; qi < 5; qi++ { + q := make([]float32, dim) + for j := range q { + q[j] = float32(rng.NormFloat64()) + } + want := original.Search(ctx, q, 10) + got := loaded.Search(ctx, q, 10) + wantURLs := make([]string, len(want)) + gotURLs := make([]string, len(got)) + for i := range want { + wantURLs[i] = want[i].URL + } + for i := range got { + gotURLs[i] = got[i].URL + } + if !sameURLOrder(gotURLs, wantURLs) { + t.Errorf("query %d: windowed round-trip diverged\n want: %v\n got: %v", qi, wantURLs, gotURLs) + } + } +} + +// TestHNSWPersistFailureLeavesNoMeta pins the meta-written-last invariant: a +// persist killed between windows must leave the store loadable as "absent", +// never with meta pointing past the written nodes. +func TestHNSWPersistFailureLeavesNoMeta(t *testing.T) { + dir := filepath.Join(t.TempDir(), "pebble") + ps, err := store.OpenPebble(dir) + if err != nil { + t.Fatalf("OpenPebble: %v", err) + } + defer ps.Close() + + const ( + n = 500 + dim = 32 + ) + h := buildTestHNSW(n, dim, 13, 7) + + savedWindow := persistWindowBytes + persistWindowBytes = 16 * 1024 + failNow := false + persistFlushed = func(nodes, bytes int) { failNow = true } + defer func() { + persistWindowBytes = savedWindow + persistFlushed = nil + }() + + ctx := &failAfterFlagCtx{Context: context.Background(), fail: &failNow} + if err := h.Persist(ctx, ps); err == nil { + t.Fatal("persist: want error from cancelled second window, got nil") + } + + loaded, ok, err := LoadHNSW(context.Background(), ps) + if err != nil { + t.Fatalf("load after failed persist: %v", err) + } + if ok || loaded != nil { + t.Fatalf("load after failed persist: want absent (ok=false), got ok=%v", ok) + } +} + +// TestHNSWPersistWindowOne: degenerate one-blob windows must still cover +// every node and load back correctly (also exercises the empty final flush). +func TestHNSWPersistWindowOne(t *testing.T) { + dir := filepath.Join(t.TempDir(), "pebble") + ps, err := store.OpenPebble(dir) + if err != nil { + t.Fatalf("OpenPebble: %v", err) + } + defer ps.Close() + ctx := context.Background() + + const n = 50 + h := buildTestHNSW(n, 32, 13, 7) + + savedWindow := persistWindowBytes + persistWindowBytes = 1 + windows := 0 + persistFlushed = func(nodes, bytes int) { + windows++ + if nodes != 1 { + t.Errorf("window %d: want 1 node per window, got %d", windows, nodes) + } + } + defer func() { + persistWindowBytes = savedWindow + persistFlushed = nil + }() + + if err := h.Persist(ctx, ps); err != nil { + t.Fatalf("persist: %v", err) + } + if windows != n { + t.Fatalf("windows: want %d, got %d", n, windows) + } + loaded, ok, err := LoadHNSW(ctx, ps) + if err != nil || !ok { + t.Fatalf("load: ok=%v err=%v", ok, err) + } + if loaded.Len() != n { + t.Errorf("loaded.Len(): want %d, got %d", n, loaded.Len()) + } +} + +// TestHNSWPersistCancelledBeforeWrite: a context dead before the first window +// commits must error out of the final flush and leave the store absent. +func TestHNSWPersistCancelledBeforeWrite(t *testing.T) { + dir := filepath.Join(t.TempDir(), "pebble") + ps, err := store.OpenPebble(dir) + if err != nil { + t.Fatalf("OpenPebble: %v", err) + } + defer ps.Close() + + h := buildTestHNSW(50, 32, 13, 7) + dead := true + ctx := &failAfterFlagCtx{Context: context.Background(), fail: &dead} + if err := h.Persist(ctx, ps); err == nil { + t.Fatal("persist: want error with pre-cancelled ctx, got nil") + } + _, ok, err := LoadHNSW(context.Background(), ps) + if err != nil { + t.Fatalf("load: %v", err) + } + if ok { + t.Fatal("load: want absent after fully-cancelled persist") + } +} + +// TestHNSWPersistMetaWriteFails: nodes land but the trailing meta write +// fails — the store must still read as absent (meta is the commit point). +func TestHNSWPersistMetaWriteFails(t *testing.T) { + dir := filepath.Join(t.TempDir(), "pebble") + ps, err := store.OpenPebble(dir) + if err != nil { + t.Fatalf("OpenPebble: %v", err) + } + defer ps.Close() + + h := buildTestHNSW(50, 32, 13, 7) + failNow := false + persistFlushed = func(nodes, bytes int) { failNow = true } + defer func() { persistFlushed = nil }() + + ctx := &failAfterFlagCtx{Context: context.Background(), fail: &failNow} + err = h.Persist(ctx, ps) + if err == nil { + t.Fatal("persist: want meta-write error, got nil") + } + if !strings.Contains(err.Error(), "put vector meta") { + t.Fatalf("persist error: want meta-write failure, got: %v", err) + } + _, ok, err := LoadHNSW(context.Background(), ps) + if err != nil { + t.Fatalf("load: %v", err) + } + if ok { + t.Fatal("load: want absent when meta never landed") + } +}