From 79dc7a356e30288f7ec260254ddb577d20acf650 Mon Sep 17 00:00:00 2001 From: Teodor Calin Date: Thu, 23 Jul 2026 23:55:55 +0300 Subject: [PATCH 1/2] feat(index): add HNSW writer bridge for Pebble crawl path + ops systemd units MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit internal/index/hnsw_writer.go implements PassageWriter against a Pebble store + in-memory HNSW graph, with periodic checkpointing and a manual Flush for shutdown. Covered by hnsw_writer_test.go (bridges passages into the graph, and verifies the auto-flush-at-threshold contract). Not yet wired into cmd_crawl/cmd_serve — those already use their own hnswPassageWriter in cmd/cosift; this is an alternate, independently tested implementation living in internal/index for review. Also adds the systemd unit/timer files used to run cosift-crawl, cosift-serve, cosift-snapshot, and the ollama service(s) in production. No secrets included — cosift-snapshot.service loads its admin token via EnvironmentFile from /etc/cosift/snapshot.env, not inline. Co-Authored-By: Claude Opus 4.8 --- internal/index/hnsw_writer.go | 105 +++++++++++++++++++++++ internal/index/hnsw_writer_test.go | 126 ++++++++++++++++++++++++++++ scripts/cosift-crawl.service | 22 +++++ scripts/cosift-crawl.timer | 14 ++++ scripts/cosift-serve.service | 30 +++++++ scripts/cosift-snapshot.service | 19 +++++ scripts/cosift-snapshot.timer | 13 +++ scripts/ollama-replica.service.tmpl | 18 ++++ scripts/ollama.service | 16 ++++ 9 files changed, 363 insertions(+) create mode 100644 internal/index/hnsw_writer.go create mode 100644 internal/index/hnsw_writer_test.go create mode 100644 scripts/cosift-crawl.service create mode 100644 scripts/cosift-crawl.timer create mode 100644 scripts/cosift-serve.service create mode 100644 scripts/cosift-snapshot.service create mode 100644 scripts/cosift-snapshot.timer create mode 100644 scripts/ollama-replica.service.tmpl create mode 100644 scripts/ollama.service diff --git a/internal/index/hnsw_writer.go b/internal/index/hnsw_writer.go new file mode 100644 index 0000000..53d59e5 --- /dev/null +++ b/internal/index/hnsw_writer.go @@ -0,0 +1,105 @@ +// HNSW write bridge for the Pebble-backed crawler path. +// +// The PassageWriter interface lets the crawler stream passage +// vectors to whatever backend the operator wires. On SQLite, *store.Store +// satisfies it directly (writes go to the passages table). On Pebble there +// was no implementation — vector indexing during crawl was silently +// skipped, which made the Pebble path BM25-only. +// +// HNSWWriter bridges the gap: passages flow into an in-memory HNSW graph +// and get periodically flushed to the PebbleStore via HNSW.Persist. +// Cheap per-passage cost (one graph insert + one +// GetDocMeta lookup); the heavy on-disk write is amortized across +// many inserts. + +package index + +import ( + "context" + "fmt" + "sync" + + "github.com/pilot-protocol/cosift/internal/store" +) + +// HNSWWriter satisfies the crawler's PassageWriter interface against a +// PebbleStore + HNSW pair. Passage URLs/titles are resolved from the +// 'i' family (cheap, no full Document gob decode). +// +// Persistence is checkpointed every PersistEvery passages. Setting +// PersistEvery to 0 (the default if you don't configure it) means +// "never auto-persist" — callers can call Persist manually at safe +// points (e.g. crawler shutdown). +type HNSWWriter struct { + hnsw *HNSW + store *store.PebbleStore + persistEvery int + + mu sync.Mutex + sinceLastFlush int +} + +// NewHNSWWriter returns a PassageWriter that adds each incoming passage +// to the HNSW graph and persists to the PebbleStore every persistEvery +// passages. persistEvery=0 disables auto-persist (caller drives it via +// HNSWWriter.Flush). +func NewHNSWWriter(h *HNSW, ps *store.PebbleStore, persistEvery int) *HNSWWriter { + return &HNSWWriter{hnsw: h, store: ps, persistEvery: persistEvery} +} + +// UpsertPassage satisfies crawler.PassageWriter. Resolves the doc's URL + +// title via PebbleStore.GetDocMeta, adds the passage to HNSW, and +// triggers a checkpoint Persist when the threshold is reached. +// +// store.Passage carries Embedding + DocID + Offset/Length + Model; +// HNSW.AddPassage needs URL + title for hit-display. +func (w *HNSWWriter) UpsertPassage(ctx context.Context, p *store.Passage) error { + if err := ctx.Err(); err != nil { + return err + } + url, title, ok, err := w.store.GetDocMeta(ctx, p.DocID) + if err != nil { + return fmt.Errorf("lookup doc %d: %w", p.DocID, err) + } + if !ok { + // Doc was deleted between Upsert and Passage write. Skip silently — + // no point indexing a vector pointing at nothing. + return nil + } + + w.mu.Lock() + defer w.mu.Unlock() + w.hnsw.AddPassage(url, title, p.Offset, p.Length, p.Embedding) + w.sinceLastFlush++ + if w.persistEvery > 0 && w.sinceLastFlush >= w.persistEvery { + w.sinceLastFlush = 0 + // Persist holds the HNSW RLock; we release our own writer mutex + // briefly first? Actually no — w.mu protects sinceLastFlush, not + // HNSW. HNSW.Persist takes its own RLock independent of ours. + if err := w.hnsw.Persist(ctx, w.store); err != nil { + return fmt.Errorf("hnsw persist: %w", err) + } + } + return nil +} + +// MarkURLInvalid zeros out every HNSW node whose url matches. Returns +// the count zeroed. Lets the crawler tell us "this URL's prior passages +// are stale" before pushing the new chunk batch. Satisfies the optional +// crawler.URLInvalidator interface. +func (w *HNSWWriter) MarkURLInvalid(ctx context.Context, url string) (int, error) { + if err := ctx.Err(); err != nil { + return 0, err + } + return w.hnsw.MarkURLPassagesInvalid(url), nil +} + +// Flush forces an immediate Persist of the current HNSW state. Crawler +// callers should invoke this at shutdown so the last partial batch of +// in-memory inserts isn't lost on next startup. +func (w *HNSWWriter) Flush(ctx context.Context) error { + w.mu.Lock() + defer w.mu.Unlock() + w.sinceLastFlush = 0 + return w.hnsw.Persist(ctx, w.store) +} diff --git a/internal/index/hnsw_writer_test.go b/internal/index/hnsw_writer_test.go new file mode 100644 index 0000000..56a0922 --- /dev/null +++ b/internal/index/hnsw_writer_test.go @@ -0,0 +1,126 @@ +package index + +import ( + "context" + "path/filepath" + "testing" + "time" + + "github.com/pilot-protocol/cosift/internal/store" +) + +// TestHNSWWriterBridgesPassagesIntoGraph — The PassageWriter +// satisfies the crawler interface and routes incoming passages into the +// HNSW graph + the PebbleStore. +func TestHNSWWriterBridgesPassagesIntoGraph(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() + + // Need a Document in the store so GetDocMeta resolves URL+title. + ctx := context.Background() + id, err := ps.UpsertDocument(ctx, &store.Document{ + URL: "https://x/doc1", Title: "Doc One", FetchedAt: time.Now(), + }) + if err != nil { + t.Fatalf("upsert: %v", err) + } + + h := NewHNSW(4) + w := NewHNSWWriter(h, ps, 0) // no auto-flush; test drives Flush manually + + vec := []float32{0.7, 0.1, 0.1, 0.7} + if err := w.UpsertPassage(ctx, &store.Passage{ + DocID: id, Offset: 0, Length: 50, Model: "test", Embedding: vec, + }); err != nil { + t.Fatalf("upsert passage: %v", err) + } + + // HNSW now holds one entry; search returns the right URL. + hits := h.Search(ctx, vec, 5) + if len(hits) != 1 { + t.Fatalf("want 1 hit, got %d", len(hits)) + } + if hits[0].URL != "https://x/doc1" { + t.Errorf("hit URL: want https://x/doc1, got %s", hits[0].URL) + } + if hits[0].Title != "Doc One" { + t.Errorf("hit Title: got %q", hits[0].Title) + } + + // Persist + reload — confirm the bridge produces a recoverable graph. + if err := w.Flush(ctx); err != nil { + t.Fatalf("flush: %v", err) + } + loaded, ok, err := LoadHNSW(ctx, ps) + if err != nil || !ok { + t.Fatalf("load: ok=%v err=%v", ok, err) + } + loadedHits := loaded.Search(ctx, vec, 5) + if len(loadedHits) != 1 || loadedHits[0].URL != "https://x/doc1" { + t.Errorf("post-reload hits: %+v", loadedHits) + } +} + +// TestHNSWWriterAutoFlushOnThreshold — when persistEvery is set, the +// writer triggers Persist automatically after N inserts. Locks in the +// auto-checkpoint contract so callers don't have to poll Flush manually. +func TestHNSWWriterAutoFlushOnThreshold(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() + // Three docs to write three passages for. + for i := 1; i <= 3; i++ { + _, err := ps.UpsertDocument(ctx, &store.Document{ + URL: "https://x/" + string(rune('a'+i-1)), + Title: "doc", + FetchedAt: time.Now(), + }) + if err != nil { + t.Fatalf("upsert %d: %v", i, err) + } + } + + h := NewHNSW(4) + w := NewHNSWWriter(h, ps, 2) // auto-persist every 2 passages + + vec := []float32{1, 0, 0, 0} + for i := int64(1); i <= 3; i++ { + if err := w.UpsertPassage(ctx, &store.Passage{ + DocID: i, Offset: 0, Length: 50, Model: "test", Embedding: vec, + }); err != nil { + t.Fatalf("passage %d: %v", i, err) + } + } + + // After 3 passages and persistEvery=2, exactly one auto-persist + // happened (after the 2nd). The 3rd is in memory but not yet on disk — + // load should see 2 nodes, not 3. + loaded, ok, err := LoadHNSW(ctx, ps) + if err != nil || !ok { + t.Fatalf("load: ok=%v err=%v", ok, err) + } + // Loaded must have at least the auto-flushed batch; in-memory writer + // has all three. After explicit Flush, persisted count catches up. + if loaded.Len() < 2 { + t.Errorf("auto-flush at threshold 2: want loaded ≥2 nodes, got %d", loaded.Len()) + } + if err := w.Flush(ctx); err != nil { + t.Fatalf("manual flush: %v", err) + } + loaded2, ok, err := LoadHNSW(ctx, ps) + if err != nil || !ok { + t.Fatalf("re-load: ok=%v err=%v", ok, err) + } + if loaded2.Len() != 3 { + t.Errorf("post-flush count: want 3, got %d", loaded2.Len()) + } +} diff --git a/scripts/cosift-crawl.service b/scripts/cosift-crawl.service new file mode 100644 index 0000000..5ef3af0 --- /dev/null +++ b/scripts/cosift-crawl.service @@ -0,0 +1,22 @@ +[Unit] +Description=cosift crawl (oneshot — stops cosift-serve, crawls, restarts serve) +# Both services want the Pebble writer lock; declare mutual exclusion so +# starting one cleanly stops the other. +Conflicts=cosift-serve.service + +[Service] +Type=oneshot +User=ubuntu +Group=ubuntu +WorkingDirectory=/home/ubuntu +Environment=COSIFT_HNSW_CHECKPOINT_SEC=60 +# Run for up to 30 minutes per cycle. Periodic HNSW checkpoint means we +# keep all vectors even if SIGTERM cuts us short. +ExecStart=/home/ubuntu/cosift -config /home/ubuntu/cosift.json crawl -backend pebble -seeds-file /home/ubuntu/seeds.txt -duration 30m +ExecStartPost=/bin/systemctl start cosift-serve.service +TimeoutStartSec=2400 +StandardOutput=journal +StandardError=journal + +[Install] +WantedBy=multi-user.target diff --git a/scripts/cosift-crawl.timer b/scripts/cosift-crawl.timer new file mode 100644 index 0000000..13dfab4 --- /dev/null +++ b/scripts/cosift-crawl.timer @@ -0,0 +1,14 @@ +[Unit] +Description=cosift crawl — cycle every 6h +# Triggers cosift-crawl.service, which stops cosift-serve.service via the +# Conflicts= directive, runs the crawl, then ExecStartPost restarts serve. + +[Timer] +# 4 cycles per day. Crawl runs 30 min, serve runs the other ~85 min before +# the next trigger. Tune to taste. +OnCalendar=*-*-* 00,06,12,18:00:00 +Persistent=true +Unit=cosift-crawl.service + +[Install] +WantedBy=timers.target diff --git a/scripts/cosift-serve.service b/scripts/cosift-serve.service new file mode 100644 index 0000000..9b107ab --- /dev/null +++ b/scripts/cosift-serve.service @@ -0,0 +1,30 @@ +[Unit] +Description=cosift pebble-serve (with in-process continuous crawler) +After=network-online.target ollama.service +Wants=network-online.target + +[Service] +Type=simple +User=ubuntu +Group=ubuntu +WorkingDirectory=/home/ubuntu +Environment=COSIFT_LOAD_HNSW=true +Environment=COSIFT_RATELIMIT_RPM=240 +Environment=COSIFT_RATELIMIT_BURST=80 +Environment=COSIFT_RATELIMIT_WHITELIST=104.28.216.88,127.0.0.1 +Environment=COSIFT_PPROF_ADDR=127.0.0.1:6060 +Environment=COSIFT_MUTEX_PROFILE_FRACTION=100 +Environment=COSIFT_DISABLE_PQ=true +Environment=COSIFT_HNSW_EF_SEARCH=200 +Environment=COSIFT_CRAWL_EMBED_CONCURRENCY=16 +Environment=COSIFT_PEBBLE_CACHE_MB=32768 +ExecStart=/home/ubuntu/cosift -config /home/ubuntu/cosift.json pebble-serve -dir /home/ubuntu/cosift-data/pebble -addr 127.0.0.1:7777 -crawl-seeds-file /home/ubuntu/seeds.txt -crawl-checkpoint 60s +Restart=on-failure +RestartSec=5 +LimitNOFILE=131072 +TimeoutStopSec=30 +MemoryMax=262G +MemoryHigh=240G + +[Install] +WantedBy=multi-user.target diff --git a/scripts/cosift-snapshot.service b/scripts/cosift-snapshot.service new file mode 100644 index 0000000..2f86294 --- /dev/null +++ b/scripts/cosift-snapshot.service @@ -0,0 +1,19 @@ +[Unit] +Description=cosift snapshot — tarball pebble + config, upload to GCS +After=network-online.target +Wants=network-online.target + +[Service] +Type=oneshot +User=ubuntu +Group=ubuntu +WorkingDirectory=/home/ubuntu +Environment=COSIFT_ADMIN_URL=http://127.0.0.1:7777 +Environment=COSIFT_CONFIG=/home/ubuntu/cosift.json +Environment=COSIFT_GCS_BUCKET=gs://pilot-cosift-index +Environment=COSIFT_KEEP=14 +# COSIFT_ADMIN_TOKEN — set via /etc/cosift/snapshot.env if cluster.peer_auth_token is configured +EnvironmentFile=-/etc/cosift/snapshot.env +ExecStart=/usr/bin/bash /home/ubuntu/snapshot.sh +StandardOutput=journal +StandardError=journal diff --git a/scripts/cosift-snapshot.timer b/scripts/cosift-snapshot.timer new file mode 100644 index 0000000..834d86d --- /dev/null +++ b/scripts/cosift-snapshot.timer @@ -0,0 +1,13 @@ +[Unit] +Description=cosift snapshot — every 4h +# Triggers cosift-snapshot.service. 4-hour cadence × KEEP=14 = ~56 hours +# of recovery history. Tune by changing OnCalendar + COSIFT_KEEP env. + +[Timer] +OnCalendar=*-*-* 00,04,08,12,16,20:00:00 +Persistent=true +RandomizedDelaySec=60 +Unit=cosift-snapshot.service + +[Install] +WantedBy=timers.target diff --git a/scripts/ollama-replica.service.tmpl b/scripts/ollama-replica.service.tmpl new file mode 100644 index 0000000..b01cc72 --- /dev/null +++ b/scripts/ollama-replica.service.tmpl @@ -0,0 +1,18 @@ +[Unit] +Description=Ollama Service (replica 1, port 11435) +After=network-online.target + +[Service] +ExecStart=/usr/local/bin/ollama serve +User=ollama +Group=ollama +Restart=always +RestartSec=3 +Environment="PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" +Environment="OLLAMA_HOST=127.0.0.1:11435" +Environment="OLLAMA_NUM_PARALLEL=32" +Environment="OLLAMA_KEEP_ALIVE=24h" +Environment="HOME=/usr/share/ollama" + +[Install] +WantedBy=default.target diff --git a/scripts/ollama.service b/scripts/ollama.service new file mode 100644 index 0000000..875e97a --- /dev/null +++ b/scripts/ollama.service @@ -0,0 +1,16 @@ +[Unit] +Description=Ollama Service +After=network-online.target + +[Service] +ExecStart=/usr/local/bin/ollama serve +Environment="OLLAMA_NUM_PARALLEL=10" +Environment="OLLAMA_KEEP_ALIVE=24h" +User=ollama +Group=ollama +Restart=always +RestartSec=3 +Environment="PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin" + +[Install] +WantedBy=default.target From d28e25a84183038023bd83652511ad7f42958406 Mon Sep 17 00:00:00 2001 From: Teodor Calin Date: Sun, 26 Jul 2026 15:07:59 +0300 Subject: [PATCH 2/2] Harden admin auth, embed-backfill durability, and the embed queue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit M23 — embed-backfill only added vectors to the in-memory HNSW graph, so a restart discarded everything the backfill produced. Snapshot the graph to the store when the handler finishes and report persist status in the response. L20 — the pebble HTTP path compared the peer/admin bearer token with a plain string comparison. Route all 24 admin gates through a shared peerTokenOK helper backed by subtle.ConstantTimeCompare. L21 — the zombie-reclaim diagnostic counter was a plain int read and written from every crawl worker. Make it an atomic.Int64. L22 — the decoupled embed queue dropped a document's passages immediately whenever the buffer was momentarily full, and its workers ran on the crawl context, which the terminator cancels as soon as the frontier drains, so queued jobs died with a cancelled context. Producers now wait a bounded interval for buffer space; embed workers run on a context detached from the crawl cancellation, with a per-job timeout and a post-crawl drain deadline so the shutdown wall clock stays bounded. Tests cover each: backfill persistence round-trips through the store, the token helper's accept/reject matrix plus a source guard, a race-detector test for the diagnostic counter, and an end-to-end test that all queued embed jobs complete after the crawl context is cancelled. Co-Authored-By: Claude Opus 5 --- cmd/cosift/feedback.go | 8 +- cmd/cosift/querylog.go | 9 +- cmd/cosift/sec_findings_test.go | 176 ++++++++++++++++++ cmd/cosift/serve_admin.go | 99 +++++----- cmd/cosift/serve_crawl.go | 126 +++++-------- cmd/cosift/serve_helpers.go | 20 +++ cmd/cosift/serve_setup.go | 9 +- internal/crawler/crawler.go | 184 +++++++++++++------ internal/crawler/embed_drain_test.go | 259 +++++++++++++++++++++++++++ 9 files changed, 687 insertions(+), 203 deletions(-) create mode 100644 cmd/cosift/sec_findings_test.go create mode 100644 internal/crawler/embed_drain_test.go diff --git a/cmd/cosift/feedback.go b/cmd/cosift/feedback.go index 16a9901..50946d4 100644 --- a/cmd/cosift/feedback.go +++ b/cmd/cosift/feedback.go @@ -108,11 +108,9 @@ func (s *pebbleHTTP) handleFeedback(w http.ResponseWriter, r *http.Request) { // reward/penalty tallies joined to the query log by qid. This is the "make it // useful" surface — the labeled signal without an external pipeline. func (s *pebbleHTTP) handleFeedbackList(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - if r.Header.Get("Authorization") != "Bearer "+want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } fp := feedbackLogPath() if fp == "" { diff --git a/cmd/cosift/querylog.go b/cmd/cosift/querylog.go index 3f3c985..f909cc3 100644 --- a/cmd/cosift/querylog.go +++ b/cmd/cosift/querylog.go @@ -86,12 +86,9 @@ func (s *pebbleHTTP) writeQueryLog(rec queryLogRec) { // handleQueryLog tails the query log. ?n=N (default 100, max 5000) returns the // last N JSON lines as a JSON array; ?raw=1 returns them as raw JSONL. func (s *pebbleHTTP) handleQueryLog(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := r.Header.Get("Authorization") - if got != "Bearer "+want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } path := os.Getenv("COSIFT_QUERY_LOG") if path == "" { diff --git a/cmd/cosift/sec_findings_test.go b/cmd/cosift/sec_findings_test.go new file mode 100644 index 0000000..7486c26 --- /dev/null +++ b/cmd/cosift/sec_findings_test.go @@ -0,0 +1,176 @@ +package main + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "regexp" + "strconv" + "strings" + "testing" + + "github.com/pilot-protocol/cosift/internal/config" + "github.com/pilot-protocol/cosift/internal/index" +) + +// --------------------------------------------------------------------------- +// peer/admin token comparison +// --------------------------------------------------------------------------- + +func TestPeerTokenOK(t *testing.T) { + req := func(auth string) *http.Request { + r := httptest.NewRequest(http.MethodGet, "/admin/x", nil) + if auth != "" { + r.Header.Set("Authorization", auth) + } + return r + } + cases := []struct { + name string + auth string + want string + ok bool + }{ + {"empty want accepts anything", "", "", true}, + {"empty want accepts a bogus token", "Bearer nonsense", "", true}, + {"exact match", "Bearer s3cret", "s3cret", true}, + {"wrong token", "Bearer s3cres", "s3cret", false}, + {"prefix of the real token", "Bearer s3c", "s3cret", false}, + {"real token plus suffix", "Bearer s3cretX", "s3cret", false}, + {"missing header", "", "s3cret", false}, + {"wrong scheme", "Basic s3cret", "s3cret", false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := peerTokenOK(req(tc.auth), tc.want); got != tc.ok { + t.Fatalf("peerTokenOK(%q, %q) = %v, want %v", tc.auth, tc.want, got, tc.ok) + } + }) + } +} + +// bearerCompare matches a direct string comparison of a request-supplied +// bearer credential against the configured token. Every admin gate must route +// through peerTokenOK instead, which compares in constant time. +var bearerCompare = regexp.MustCompile( + `(?:got != want|!= "Bearer "\+want|r\.Header\.Get\("Authorization"\) != )`) + +func TestAdminGatesUseConstantTimeCompare(t *testing.T) { + files, err := filepath.Glob("*.go") + if err != nil { + t.Fatalf("glob: %v", err) + } + var offenders []string + for _, f := range files { + if strings.HasSuffix(f, "_test.go") { + continue + } + src, err := os.ReadFile(f) + if err != nil { + t.Fatalf("read %s: %v", f, err) + } + for i, line := range strings.Split(string(src), "\n") { + if bearerCompare.MatchString(line) { + offenders = append(offenders, f+":"+strconv.Itoa(i+1)+": "+strings.TrimSpace(line)) + } + } + } + if len(offenders) > 0 { + t.Fatalf("token comparisons not routed through peerTokenOK:\n %s", + strings.Join(offenders, "\n ")) + } +} + +// requireAdmin must still reject a wrong token and admit the right one. +func TestRequireAdminGate(t *testing.T) { + s := &pebbleHTTP{} + s.cluster = config.Cluster{PeerAuthToken: "topsecret"} + called := 0 + h := s.requireAdmin(func(w http.ResponseWriter, r *http.Request) { + called++ + w.WriteHeader(http.StatusOK) + }) + + r := httptest.NewRequest(http.MethodPost, "/admin/x", nil) + r.Header.Set("Authorization", "Bearer wrong") + w := httptest.NewRecorder() + h(w, r) + if w.Code != http.StatusUnauthorized { + t.Fatalf("wrong token: status %d, want 401", w.Code) + } + if called != 0 { + t.Fatalf("handler ran for a wrong token") + } + + r = httptest.NewRequest(http.MethodPost, "/admin/x", nil) + r.Header.Set("Authorization", "Bearer topsecret") + w = httptest.NewRecorder() + h(w, r) + if w.Code != http.StatusOK || called != 1 { + t.Fatalf("right token: status %d called %d, want 200/1", w.Code, called) + } +} + +// --------------------------------------------------------------------------- +// embed-backfill durability +// --------------------------------------------------------------------------- + +// TestEmbedBackfillPersistsHNSW checks that vectors added by the backfill +// handler are written to the store, not just to the in-memory graph: after the +// handler returns, loading the graph back from Pebble must yield the nodes. +func TestEmbedBackfillPersistsHNSW(t *testing.T) { + f := populatedPebbleStore(t) + mock := openaiTestServer(t) + srv := f.makeServer(mock) + // Start from an empty graph so every corpus doc counts as missing. + srv.hnsw = index.NewHNSW(f.dim) + + ctx := context.Background() + if _, ok, err := index.LoadHNSWMeta(ctx, f.ps); err != nil { + t.Fatalf("LoadHNSWMeta: %v", err) + } else if ok { + t.Fatalf("precondition: store already holds HNSW meta") + } + + r := httptest.NewRequest(http.MethodPost, "/admin/embed-backfill", + strings.NewReader(`{"workers":2}`)) + w := httptest.NewRecorder() + srv.handleEmbedBackfill(w, r) + if w.Code != http.StatusOK { + t.Fatalf("status %d: %s", w.Code, w.Body.String()) + } + + var resp struct { + Missing int `json:"missing"` + Embedded int `json:"embedded"` + Persisted bool `json:"persisted"` + PersistError string `json:"persist_error"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode: %v (%s)", err, w.Body.String()) + } + if resp.Embedded == 0 { + t.Fatalf("nothing embedded: %s", w.Body.String()) + } + if !resp.Persisted { + t.Fatalf("handler reported persisted=false (err=%q)", resp.PersistError) + } + if srv.hnsw.Len() == 0 { + t.Fatalf("in-memory graph is empty after backfill") + } + + loaded, ok, err := index.LoadHNSW(ctx, f.ps) + if err != nil { + t.Fatalf("LoadHNSW: %v", err) + } + if !ok { + t.Fatalf("backfilled vectors were not persisted to the store") + } + if loaded.Len() != srv.hnsw.Len() { + t.Fatalf("persisted graph has %d nodes, in-memory has %d", + loaded.Len(), srv.hnsw.Len()) + } +} diff --git a/cmd/cosift/serve_admin.go b/cmd/cosift/serve_admin.go index 3eabfa8..ddc814d 100644 --- a/cmd/cosift/serve_admin.go +++ b/cmd/cosift/serve_admin.go @@ -26,12 +26,9 @@ import ( // against the codebook already loaded on this serve. Order-of-magnitude // faster than handlePQTrain because it skips k-means; just encode loop. func (s *pebbleHTTP) handlePQEncode(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.hnsw == nil { writeProblem(w, http.StatusBadRequest, "pq-encode: no HNSW loaded") @@ -90,12 +87,9 @@ func (s *pebbleHTTP) handlePQEncode(w http.ResponseWriter, r *http.Request) { // returned path is safe to tar — Pebble's compactor cannot mutate hard-linked // SSTs. Caller is responsible for deleting the dir after consuming it. func (s *pebbleHTTP) handleCheckpoint(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } base := os.Getenv("COSIFT_CHECKPOINT_DIR") if base == "" { @@ -127,12 +121,9 @@ type pqTrainReq struct { } func (s *pebbleHTTP) handlePQTrain(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.hnsw == nil || s.hnsw.Len() == 0 { writeProblem(w, http.StatusBadRequest, "pq-train requires an in-memory HNSW with nodes") @@ -246,12 +237,9 @@ var evalQuickQueries = []string{ // in-process (re-uses the same chat/retrieval stack as a real /answer call), // so the rate that comes back matches what users see. func (s *pebbleHTTP) handleEvalQuick(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.chat == nil { writeProblem(w, http.StatusNotImplemented, "eval-quick requires cfg.Chat.Model") @@ -361,12 +349,9 @@ func (s *pebbleHTTP) handleEvalQuick(w http.ResponseWriter, r *http.Request) { // ResponseController because compacting a multi-million-node graph routinely // runs past 60s. Returns counters so operators can confirm progress. func (s *pebbleHTTP) handleHNSWCompact(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.hnsw == nil { writeProblem(w, http.StatusNotImplemented, "hnsw-compact requires a loaded HNSW graph") @@ -474,12 +459,9 @@ type embedBackfillReq struct { // After it logs "done", set COSIFT_HOST_PARTITION_READ=1 (and =1 for the write // flag to keep it fresh) and restart. func (s *pebbleHTTP) handleHostBackfill(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.store == nil { writeProblem(w, http.StatusNotImplemented, "host-backfill requires a PebbleStore") @@ -510,12 +492,9 @@ func (s *pebbleHTTP) handleHostBackfill(w http.ResponseWriter, r *http.Request) } func (s *pebbleHTTP) handleEmbedBackfill(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.embedder == nil || s.hnsw == nil { writeProblem(w, http.StatusNotImplemented, "embed backfill requires both an embedder and HNSW graph to be configured") @@ -599,14 +578,34 @@ func (s *pebbleHTTP) handleEmbedBackfill(w http.ResponseWriter, r *http.Request) } wg.Wait() - log.Printf("embed-backfill: scanned=%d missing=%d embedded=%d in %s", - scanned.Load(), missing.Load(), embedded.Load(), time.Since(t0).Round(time.Second)) - writeJSON(w, http.StatusOK, map[string]any{ - "scanned": scanned.Load(), - "missing": missing.Load(), - "embedded": embedded.Load(), - "elapsed": time.Since(t0).String(), - }) + // AddPassage only mutates the in-memory graph. Snapshot it to the store + // so the newly-embedded vectors survive a restart; without this the + // backfill's work is lost whenever the process exits. Detached context + // so a client disconnect can't abort the write. + persisted := false + var persistErr string + if s.store != nil && embedded.Load() > 0 { + if err := s.hnsw.Persist(context.Background(), s.store); err != nil { + persistErr = err.Error() + log.Printf("embed-backfill: hnsw persist FAILED: %v", err) + } else { + persisted = true + } + } + + log.Printf("embed-backfill: scanned=%d missing=%d embedded=%d persisted=%v in %s", + scanned.Load(), missing.Load(), embedded.Load(), persisted, time.Since(t0).Round(time.Second)) + resp := map[string]any{ + "scanned": scanned.Load(), + "missing": missing.Load(), + "embedded": embedded.Load(), + "persisted": persisted, + "elapsed": time.Since(t0).String(), + } + if persistErr != "" { + resp["persist_error"] = persistErr + } + writeJSON(w, http.StatusOK, resp) } // truncateForEmbedLite mirrors the crawler's helper without pulling the diff --git a/cmd/cosift/serve_crawl.go b/cmd/cosift/serve_crawl.go index 5039f13..bbcbdad 100644 --- a/cmd/cosift/serve_crawl.go +++ b/cmd/cosift/serve_crawl.go @@ -25,12 +25,9 @@ type crawlEnqueueReq struct { func (s *pebbleHTTP) handleCrawlEnqueue(w http.ResponseWriter, r *http.Request) { // Auth - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") + return } if s.crawlSeed == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -55,12 +52,9 @@ func (s *pebbleHTTP) handleCrawlEnqueue(w http.ResponseWriter, r *http.Request) // link-target domain recurs above their frequency threshold. Persisted so it // survives restart. Body: {"domain":"example.com"}. func (s *pebbleHTTP) handleAllowDomain(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") + return } if s.crawlAllowDomain == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler") @@ -91,12 +85,9 @@ func (s *pebbleHTTP) handleAllowDomain(w http.ResponseWriter, r *http.Request) { // approximate count of deleted entries (DeleteRange is O(tombstones), // not O(N), so the count is reported as -1 to signal "swept range"). func (s *pebbleHTTP) handleFrontierClear(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") + return } if err := s.store.ClearFrontier(r.Context()); err != nil { writeProblem(w, http.StatusInternalServerError, err.Error()) @@ -115,12 +106,9 @@ type frontierPurgeReq struct { } func (s *pebbleHTTP) handleFrontierPurgeHost(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") + return } var req frontierPurgeReq body, _ := io.ReadAll(io.LimitReader(r.Body, 64<<10)) @@ -145,12 +133,9 @@ type sitemapImportReq struct { } func (s *pebbleHTTP) handleSitemapImport(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlSeedSitemap == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -182,12 +167,9 @@ func (s *pebbleHTTP) handleSitemapImport(w http.ResponseWriter, r *http.Request) // entries are stuck in errored state — sitemap-import's INSERT OR IGNORE won't // re-queue them, but Recrawl will. func (s *pebbleHTTP) handleRecrawlSitemap(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlRecrawl == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler") @@ -252,12 +234,9 @@ type crawlNowReq struct { } func (s *pebbleHTTP) handleCrawlNow(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlFetchNow == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -310,12 +289,9 @@ type sitePackReq struct { } func (s *pebbleHTTP) handleSitePack(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlSeedSitemap == nil || s.crawlSeedRSS == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -491,12 +467,9 @@ type siteSubmitReq struct { } func (s *pebbleHTTP) handleSiteSubmit(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlSeedSitemapLane == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -572,12 +545,9 @@ type wetImportBulkReq struct { } func (s *pebbleHTTP) handleWETImportBulk(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlSeedWET == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -700,12 +670,9 @@ type wetImportReq struct { } func (s *pebbleHTTP) handleWETImport(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlSeedWET == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -744,12 +711,9 @@ type frontierDemoteHostReq struct { } func (s *pebbleHTTP) handleFrontierDemoteHost(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } ps, ok := any(s.store).(*store.PebbleStore) if !ok { @@ -788,12 +752,9 @@ func (s *pebbleHTTP) handleFrontierDemoteHost(w http.ResponseWriter, r *http.Req // key. GetLaneStats then reported impossibly-high in_flight counts // (>max_concurrent). Idempotent — re-running is a no-op once clean. func (s *pebbleHTTP) handleFrontierPurgeStaleInFlight(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } ps, ok := any(s.store).(*store.PebbleStore) if !ok { @@ -819,12 +780,9 @@ type rssImportReq struct { } func (s *pebbleHTTP) handleRSSImport(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlSeedRSS == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") diff --git a/cmd/cosift/serve_helpers.go b/cmd/cosift/serve_helpers.go index 38f38ac..e7a3747 100644 --- a/cmd/cosift/serve_helpers.go +++ b/cmd/cosift/serve_helpers.go @@ -2,6 +2,7 @@ package main import ( "context" + "crypto/subtle" _ "embed" "encoding/json" "flag" @@ -19,6 +20,25 @@ import ( "github.com/pilot-protocol/cosift/internal/store" ) +// bearerToken extracts the credential from an "Authorization: Bearer " +// request header. Returns "" when the header is absent or carries a different +// scheme. +func bearerToken(r *http.Request) string { + return strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") +} + +// peerTokenOK reports whether the request presents the expected peer/admin +// token. The comparison runs in time independent of how many leading bytes +// match, and returns 1 only when both lengths and all bytes are equal. +// An empty want means the endpoint is unauthenticated (single-node default) +// and every caller is accepted. +func peerTokenOK(r *http.Request, want string) bool { + if want == "" { + return true + } + return subtle.ConstantTimeCompare([]byte(bearerToken(r)), []byte(want)) == 1 +} + // retrievalFilters bundles the four post-retrieval predicates /search, // /answer, /find_similar, and /research all share. Centralized in // so /research's sync + stream paths stay in lockstep. diff --git a/cmd/cosift/serve_setup.go b/cmd/cosift/serve_setup.go index ece6188..b3ea893 100644 --- a/cmd/cosift/serve_setup.go +++ b/cmd/cosift/serve_setup.go @@ -1281,12 +1281,9 @@ func stripPort(remoteAddr string) string { // without auth, even if it forgets the per-handler check. func (s *pebbleHTTP) requireAdmin(h http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") + return } h(w, r) } diff --git a/internal/crawler/crawler.go b/internal/crawler/crawler.go index c9709ec..f397880 100644 --- a/internal/crawler/crawler.go +++ b/internal/crawler/crawler.go @@ -65,9 +65,10 @@ type Crawler struct { // bounded diagnostic counter for the zombie-reclaim path. // Logs the first few zero-hit cases so we can confirm the code path - // is actually reached even when no prior passages exist. Cheap (just - // an int read on the hot path; not atomic — diagnostic only). - zombieDebugLogged int + // is actually reached even when no prior passages exist. Written from + // every crawl worker goroutine, so the read-modify-write goes through + // atomic.Int64. + zombieDebugLogged atomic.Int64 // per-host error-rate tracking via sync.Map + atomic // counters. With 512 workers we cannot afford a single write lock @@ -93,6 +94,13 @@ type Crawler struct { embedDone atomic.Int64 embedFailed atomic.Int64 + // embedDrainUntil is a unix-nano wall-clock deadline for the post-crawl + // drain of embedQ. Zero means "no deadline" (crawl still running). Run + // sets it just before closing embedQ; embed workers stop doing work and + // count the remainder as dropped once it passes, so shutdown stays + // bounded even when the embedder is slow or unreachable. + embedDrainUntil atomic.Int64 + // URLs whose frontier transition happened inside processClaimed via // WriteCrawlResult. The worker checks this set before its own // CompleteFrontier call so we don't pay a redundant mu hop. @@ -587,9 +595,15 @@ func (c *Crawler) Run(ctx context.Context) error { if embedWorkers > 0 { bufSize := envIntCrawler("COSIFT_EMBED_DECOUPLE_BUFFER", 4096) c.embedQ = make(chan *embedJob, bufSize) + // runCtx is cancelled the moment the frontier drains (see + // terminator), while jobs already queued still need a live + // context to embed and write. Detach the embed pool from that + // cancellation; its wall-clock bound comes from embedDrainUntil + // plus the per-job timeout instead. + embedCtx := context.WithoutCancel(ctx) for i := 0; i < embedWorkers; i++ { embedWG.Add(1) - go c.embedWorker(runCtx, &embedWG) + go c.embedWorker(embedCtx, &embedWG) } log.Printf("crawler: embed decouple ON (%d workers, %d-buf)", embedWorkers, bufSize) } @@ -620,8 +634,14 @@ func (c *Crawler) Run(ctx context.Context) error { // then wait for embed workers to drain. Otherwise an early close // would race a still-running crawl worker's send and panic. if c.embedQ != nil { + // Cap the drain so a slow or unreachable embedder can't hold the + // crawl open indefinitely. Jobs still queued past the deadline are + // counted as dropped and left for embed-backfill. + drainSecs := envIntCrawler("COSIFT_EMBED_DRAIN_SECONDS", defaultEmbedDrainSeconds) + c.embedDrainUntil.Store(time.Now().Add(time.Duration(drainSecs) * time.Second).UnixNano()) close(c.embedQ) embedWG.Wait() + c.embedDrainUntil.Store(0) log.Printf("crawler: embed pool drained — queued=%d done=%d failed=%d dropped=%d", c.embedQueued.Load(), c.embedDone.Load(), c.embedFailed.Load(), c.embedDropped.Load()) c.embedQ = nil // nil so FetchAndIndexNow doesn't send to closed channel between restarts @@ -644,64 +664,126 @@ func envIntCrawler(key string, def int) int { return n } +// Defaults for the decoupled embed pipeline's wall-clock bounds. Both are +// overridable via env (COSIFT_EMBED_DRAIN_SECONDS, COSIFT_EMBED_JOB_SECONDS, +// COSIFT_EMBED_ENQUEUE_WAIT_MS). +const ( + defaultEmbedDrainSeconds = 300 + defaultEmbedJobSeconds = 120 + defaultEmbedEnqueueWaitMs = 5000 +) + +// enqueueEmbedJob hands a job to the embed pool. It first tries a +// non-blocking send, then waits a bounded amount of time for buffer space +// before giving up. Reports whether the job was accepted; a rejected job's +// passages are left for embed-backfill. +func (c *Crawler) enqueueEmbedJob(ctx context.Context, job *embedJob) bool { + select { + case c.embedQ <- job: + return true + default: + } + wait := time.Duration(envIntCrawler("COSIFT_EMBED_ENQUEUE_WAIT_MS", defaultEmbedEnqueueWaitMs)) * time.Millisecond + if wait <= 0 { + return false + } + t := time.NewTimer(wait) + defer t.Stop() + select { + case c.embedQ <- job: + return true + case <-ctx.Done(): + return false + case <-t.C: + return false + } +} + +// embedJobExpired reports whether the post-crawl drain deadline has passed. +func (c *Crawler) embedJobExpired() bool { + d := c.embedDrainUntil.Load() + return d != 0 && time.Now().UnixNano() > d +} + // embedWorker drains c.embedQ. For each job, embeds the chunk texts and // writes passages to the configured PassageWriter. Uses the optional // batch writer when available for one HNSW lock per doc instead of // one per chunk. +// +// ctx is detached from the crawl's cancellation so queued jobs still +// complete after the frontier drains; each job carries its own timeout and +// the pool stops working once the drain deadline passes. func (c *Crawler) embedWorker(ctx context.Context, wg *sync.WaitGroup) { defer wg.Done() + jobTimeout := time.Duration(envIntCrawler("COSIFT_EMBED_JOB_SECONDS", defaultEmbedJobSeconds)) * time.Second + if jobTimeout <= 0 { + jobTimeout = defaultEmbedJobSeconds * time.Second + } for job := range c.embedQ { - vecs, err := c.embedder.Embed(ctx, job.texts) - if err != nil { - c.embedFailed.Add(1) - log.Printf("embed-decouple %s: %v", job.url, err) - continue - } - if len(vecs) != len(job.chunks) { - c.embedFailed.Add(1) + if c.embedJobExpired() { + c.embedDropped.Add(1) continue } - // 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 inv, ok := c.passageWriter.(URLInvalidator); ok { - _, _ = inv.MarkURLInvalid(ctx, job.url) + c.runEmbedJob(ctx, job, jobTimeout) + } +} + +// runEmbedJob embeds one job's chunk texts and writes the resulting passages. +// The per-job context bounds how long a single document may hold a worker. +func (c *Crawler) runEmbedJob(parent context.Context, job *embedJob, jobTimeout time.Duration) { + ctx, cancel := context.WithTimeout(parent, jobTimeout) + defer cancel() + + vecs, err := c.embedder.Embed(ctx, job.texts) + if err != nil { + c.embedFailed.Add(1) + log.Printf("embed-decouple %s: %v", job.url, err) + return + } + if len(vecs) != len(job.chunks) { + c.embedFailed.Add(1) + return + } + // 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 inv, ok := c.passageWriter.(URLInvalidator); ok { + _, _ = inv.MarkURLInvalid(ctx, job.url) + } + } + // Prefer the batch interface (single HNSW lock for the whole + // doc) over per-chunk writes. + if bw, ok := c.passageWriter.(PassageWriterBatch); ok { + ps := make([]*store.Passage, len(job.chunks)) + for i, ch := range job.chunks { + ps[i] = &store.Passage{ + DocID: job.docID, + Offset: ch.Offset, + Length: ch.Length, + Model: c.embedder.Model(), + Embedding: vecs[i], } } - // Prefer the batch interface (single HNSW lock for the whole - // doc) over per-chunk writes. - if bw, ok := c.passageWriter.(PassageWriterBatch); ok { - ps := make([]*store.Passage, len(job.chunks)) - for i, ch := range job.chunks { - ps[i] = &store.Passage{ - DocID: job.docID, - Offset: ch.Offset, - Length: ch.Length, - Model: c.embedder.Model(), - Embedding: vecs[i], - } - } - if err := bw.UpsertPassageBatch(ctx, ps); err != nil { - c.embedFailed.Add(1) - log.Printf("embed-decouple batch %s: %v", job.url, err) - continue + if err := bw.UpsertPassageBatch(ctx, ps); err != nil { + c.embedFailed.Add(1) + log.Printf("embed-decouple batch %s: %v", job.url, err) + return + } + } else { + for i, ch := range job.chunks { + p := &store.Passage{ + DocID: job.docID, + Offset: ch.Offset, + Length: ch.Length, + Model: c.embedder.Model(), + Embedding: vecs[i], } - } else { - for i, ch := range job.chunks { - p := &store.Passage{ - DocID: job.docID, - Offset: ch.Offset, - Length: ch.Length, - Model: c.embedder.Model(), - Embedding: vecs[i], - } - if err := c.passageWriter.UpsertPassage(ctx, p); err != nil { - log.Printf("embed-decouple passage %s offset=%d: %v", job.url, ch.Offset, err) - } + if err := c.passageWriter.UpsertPassage(ctx, p); err != nil { + log.Printf("embed-decouple passage %s offset=%d: %v", job.url, ch.Offset, err) } } - c.embedDone.Add(1) } + c.embedDone.Add(1) } // statusDumper writes a JSON snapshot of crawl progress every 10s to path. @@ -1286,10 +1368,9 @@ func (c *Crawler) processClaimed(ctx context.Context, item store.FrontierItem, g } if c.embedQ != nil { job := &embedJob{url: item.URL, docID: id, chunks: chunks, texts: texts} - select { - case c.embedQ <- job: + if c.enqueueEmbedJob(ctx, job) { c.embedQueued.Add(1) - default: + } else { c.embedDropped.Add(1) } c.enqueueLinks(ctx, parsed.Links, item.Depth+1) @@ -1328,10 +1409,9 @@ func (c *Crawler) processClaimed(ctx context.Context, item store.FrontierItem, g log.Printf("zombie-reclaim %s: %v", item.URL, err) case n > 0: log.Printf("zombie-reclaim %s: invalidated %d stale passages", item.URL, n) - case c.zombieDebugLogged < 5: + case c.zombieDebugLogged.Add(1) <= 5: // Diagnostic: first few zero-hit cases to confirm code path is reached. log.Printf("zombie-reclaim %s: 0 prior passages (diagnostic)", item.URL) - c.zombieDebugLogged++ } } } diff --git a/internal/crawler/embed_drain_test.go b/internal/crawler/embed_drain_test.go new file mode 100644 index 0000000..32ddc2a --- /dev/null +++ b/internal/crawler/embed_drain_test.go @@ -0,0 +1,259 @@ +package crawler + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/pilot-protocol/cosift/internal/config" + "github.com/pilot-protocol/cosift/internal/store" +) + +// slowEmbedder takes a fixed amount of time per call and honors context +// cancellation the way a real HTTP-backed embedder does. +type slowEmbedder struct { + dim int + delay time.Duration + + mu sync.Mutex + calls int + failed int +} + +func (s *slowEmbedder) Model() string { return "slow-emb" } +func (s *slowEmbedder) Dim() int { return s.dim } + +func (s *slowEmbedder) Embed(ctx context.Context, texts []string) ([][]float32, error) { + t := time.NewTimer(s.delay) + defer t.Stop() + select { + case <-ctx.Done(): + s.mu.Lock() + s.failed++ + s.mu.Unlock() + return nil, ctx.Err() + case <-t.C: + } + s.mu.Lock() + s.calls++ + s.mu.Unlock() + out := make([][]float32, len(texts)) + for i := range texts { + v := make([]float32, s.dim) + for j := range v { + v[j] = 0.1 + } + out[i] = v + } + return out, nil +} + +// countingWriter records passage writes and reports zero prior passages from +// MarkURLInvalid, which is the branch that drives the zombie-reclaim +// diagnostic counter. +type countingWriter struct { + inner PassageWriter + // invalidateDelay holds each worker inside the reclaim block so that all + // of them reach the diagnostic counter at roughly the same instant. + invalidateDelay time.Duration + + mu sync.Mutex + upserts int + invalids int +} + +func (c *countingWriter) UpsertPassage(ctx context.Context, p *store.Passage) error { + if err := c.inner.UpsertPassage(ctx, p); err != nil { + return err + } + c.mu.Lock() + c.upserts++ + c.mu.Unlock() + return nil +} + +func (c *countingWriter) MarkURLInvalid(ctx context.Context, url string) (int, error) { + if c.invalidateDelay > 0 { + time.Sleep(c.invalidateDelay) + } + c.mu.Lock() + c.invalids++ + c.mu.Unlock() + return 0, nil +} + +func (c *countingWriter) counts() (int, int) { + c.mu.Lock() + defer c.mu.Unlock() + return c.upserts, c.invalids +} + +// pagesServer serves n distinct HTML pages under /p. +func pagesServer(t *testing.T, n int) (*httptest.Server, []string) { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + fmt.Fprintf(w, `Page %s +

Heading %s

Body text for path %s about distributed indexing and retrieval.

`, + r.URL.Path, r.URL.Path, r.URL.Path) + })) + t.Cleanup(srv.Close) + urls := make([]string, 0, n) + for i := 0; i < n; i++ { + urls = append(urls, fmt.Sprintf("%s/p%d", srv.URL, i)) + } + return srv, urls +} + +func fastCrawlCfg() config.Crawler { + cfg := config.Default().Crawler + cfg.MaxDepth = 0 + cfg.PerHostDelayMs = 0 + cfg.MaxConcurrent = 8 + cfg.RespectRobots = false + return cfg +} + +// TestEmbedPoolFinishesQueuedJobsAfterCrawlContextCancel checks that the +// decoupled embed pool keeps a usable context after the crawl's own context is +// cancelled (the terminator cancels it as soon as the frontier drains). Jobs +// already sitting in the queue must still be embedded and written, not failed. +func TestEmbedPoolFinishesQueuedJobsAfterCrawlContextCancel(t *testing.T) { + const pages = 40 + + t.Setenv("COSIFT_EMBED_DECOUPLE_WORKERS", "1") + t.Setenv("COSIFT_EMBED_DECOUPLE_BUFFER", "256") + + _, urls := pagesServer(t, pages) + s := newStoreT(t) + emb := &slowEmbedder{dim: 8, delay: 100 * time.Millisecond} + cw := &countingWriter{inner: s} + + c := New(fastCrawlCfg(), s).WithEmbedder(emb).WithPassageWriter(cw) + for _, u := range urls { + if err := c.Seed(u); err != nil { + t.Fatalf("seed %s: %v", u, err) + } + } + if err := c.Run(context.Background()); err != nil { + t.Fatalf("run: %v", err) + } + + queued := c.embedQueued.Load() + done := c.embedDone.Load() + failed := c.embedFailed.Load() + dropped := c.embedDropped.Load() + + if queued == 0 { + t.Fatalf("no embed jobs were queued — decoupled path did not engage") + } + if failed != 0 { + t.Errorf("embed jobs failed after the crawl context was cancelled: failed=%d (queued=%d done=%d)", + failed, queued, done) + } + if done != queued { + t.Errorf("embed pool completed %d of %d queued jobs (failed=%d dropped=%d)", + done, queued, failed, dropped) + } + if up, _ := cw.counts(); up == 0 { + t.Errorf("no passages were written") + } +} + +// TestEnqueueEmbedJobWaitsForCapacity checks that a momentarily full queue +// makes the producer wait for space rather than discarding the document's +// passages outright. +func TestEnqueueEmbedJobWaitsForCapacity(t *testing.T) { + c := newBare(config.Default().Crawler) + c.embedQ = make(chan *embedJob, 1) + c.embedQ <- &embedJob{url: "https://x.example/filler"} + + t.Setenv("COSIFT_EMBED_ENQUEUE_WAIT_MS", "2000") + + go func() { + time.Sleep(50 * time.Millisecond) + <-c.embedQ + }() + + start := time.Now() + if !c.enqueueEmbedJob(context.Background(), &embedJob{url: "https://x.example/real"}) { + t.Fatalf("job rejected even though capacity freed up after %s", time.Since(start)) + } + if elapsed := time.Since(start); elapsed < 40*time.Millisecond { + t.Fatalf("returned in %s — did not actually wait for capacity", elapsed) + } +} + +// TestEnqueueEmbedJobGivesUpWhenQueueStaysFull keeps the wall clock bounded: +// a permanently full queue must not block a crawl worker indefinitely. +func TestEnqueueEmbedJobGivesUpWhenQueueStaysFull(t *testing.T) { + c := newBare(config.Default().Crawler) + c.embedQ = make(chan *embedJob, 1) + c.embedQ <- &embedJob{url: "https://x.example/filler"} + + t.Setenv("COSIFT_EMBED_ENQUEUE_WAIT_MS", "50") + + start := time.Now() + if c.enqueueEmbedJob(context.Background(), &embedJob{url: "https://x.example/real"}) { + t.Fatalf("job accepted into a full queue") + } + if elapsed := time.Since(start); elapsed > 2*time.Second { + t.Fatalf("enqueue took %s — wait budget not honored", elapsed) + } +} + +// TestEmbedJobExpiredHonoursDrainDeadline covers the shutdown bound: once the +// drain deadline has passed, remaining queued work is abandoned. +func TestEmbedJobExpiredHonoursDrainDeadline(t *testing.T) { + c := newBare(config.Default().Crawler) + if c.embedJobExpired() { + t.Fatalf("expired with no deadline set") + } + c.embedDrainUntil.Store(time.Now().Add(time.Hour).UnixNano()) + if c.embedJobExpired() { + t.Fatalf("expired with an hour of budget left") + } + c.embedDrainUntil.Store(time.Now().Add(-time.Second).UnixNano()) + if !c.embedJobExpired() { + t.Fatalf("did not expire past the deadline") + } +} + +// TestZombieReclaimDiagnosticCounterConcurrent drives the zombie-reclaim +// diagnostic path from many crawl workers at once. Run under -race. +func TestZombieReclaimDiagnosticCounterConcurrent(t *testing.T) { + const pages = 60 + + t.Setenv("COSIFT_ZOMBIE_RECLAIM", "1") + t.Setenv("COSIFT_EMBED_DECOUPLE_WORKERS", "0") + + _, urls := pagesServer(t, pages) + s := newStoreT(t) + // slowEmbedder with no delay: the package's stubEmbedder keeps an + // unsynchronized call counter and is only safe at MaxConcurrent=1. + emb := &slowEmbedder{dim: 8} + cw := &countingWriter{inner: s, invalidateDelay: 15 * time.Millisecond} + + cfg := fastCrawlCfg() + cfg.MaxConcurrent = 16 + c := New(cfg, s).WithEmbedder(emb).WithPassageWriter(cw) + for _, u := range urls { + if err := c.Seed(u); err != nil { + t.Fatalf("seed %s: %v", u, err) + } + } + if err := c.Run(context.Background()); err != nil { + t.Fatalf("run: %v", err) + } + + if _, inv := cw.counts(); inv == 0 { + t.Fatalf("zombie-reclaim path never ran — counter was not exercised") + } + if got := c.zombieDebugLogged.Load(); got == 0 { + t.Fatalf("diagnostic counter never incremented") + } +}