From ca174e23fd492a1dd027cd9be1a25a69f5ff7e13 Mon Sep 17 00:00:00 2001 From: Andrew Mattie Date: Fri, 7 Aug 2026 13:38:19 -0500 Subject: [PATCH] fix(search): repair shards after recovered faults --- search/shard_fault_linux_test.go | 99 ++++++ search/shard_repair.go | 135 ++++++++ search/shard_repair_e2e_linux_test.go | 144 +++++++++ search/shard_repair_test.go | 434 ++++++++++++++++++++++++++ search/shards.go | 108 ++++++- 5 files changed, 904 insertions(+), 16 deletions(-) create mode 100644 search/shard_fault_linux_test.go create mode 100644 search/shard_repair.go create mode 100644 search/shard_repair_e2e_linux_test.go create mode 100644 search/shard_repair_test.go diff --git a/search/shard_fault_linux_test.go b/search/shard_fault_linux_test.go new file mode 100644 index 000000000..162bd6a29 --- /dev/null +++ b/search/shard_fault_linux_test.go @@ -0,0 +1,99 @@ +//go:build linux + +package search + +import ( + "context" + "runtime/debug" + "sync/atomic" + "testing" + + "github.com/sourcegraph/zoekt" + "github.com/sourcegraph/zoekt/query" + "golang.org/x/sys/unix" +) + +var faultReadSink byte + +type mmapFaultSearcher struct { + page []byte + faultList atomic.Bool +} + +func (s *mmapFaultSearcher) Search(context.Context, query.Q, *zoekt.SearchOptions) (*zoekt.SearchResult, error) { + faultReadSink = s.page[0] + return &zoekt.SearchResult{}, nil +} + +func (s *mmapFaultSearcher) List(context.Context, query.Q, *zoekt.ListOptions) (*zoekt.RepoList, error) { + if s.faultList.Load() { + faultReadSink = s.page[0] + } + return &zoekt.RepoList{}, nil +} + +func (s *mmapFaultSearcher) Stats() (*zoekt.RepoStats, error) { + return &zoekt.RepoStats{}, nil +} + +func (s *mmapFaultSearcher) Close() {} + +func (s *mmapFaultSearcher) String() string { return "mmap-fault.zoekt" } + +func inaccessibleMappedPage(t *testing.T) []byte { + t.Helper() + data, err := unix.Mmap( + -1, + 0, + unix.Getpagesize(), + unix.PROT_READ|unix.PROT_WRITE, + unix.MAP_PRIVATE|unix.MAP_ANONYMOUS, + ) + if err != nil { + t.Fatalf("mmap: %v", err) + } + t.Cleanup(func() { + if err := unix.Munmap(data); err != nil { + t.Errorf("munmap: %v", err) + } + }) + if err := unix.Mprotect(data, unix.PROT_NONE); err != nil { + t.Fatalf("mprotect: %v", err) + } + return data +} + +func TestShardMmapFaultsAreRecoveredAndPanicModeIsRestored(t *testing.T) { + before := debug.SetPanicOnFault(false) + defer debug.SetPanicOnFault(before) + + ss := newShardedSearcher(1) + ss.shardRepairs = newShardRepairQueue(func(string, zoekt.Searcher) error { return nil }) + searcher := &mmapFaultSearcher{page: inaccessibleMappedPage(t)} + ss.replace(map[string]zoekt.Searcher{searcher.String(): searcher}) + shard := ss.getLoaded().shards[0] + + result, err := ss.searchOneShard( + context.Background(), + shard, + &query.Const{Value: true}, + &zoekt.SearchOptions{}, + ) + if err != nil { + t.Fatalf("search returned an error: %v", err) + } + if result.Stats.Crashes != 1 { + t.Fatalf("search crashes = %d, want 1", result.Stats.Crashes) + } + + searcher.faultList.Store(true) + listResults := make(chan shardListResult, 1) + ss.listOneShard(context.Background(), shard, &query.Const{Value: true}, nil, listResults) + if result := <-listResults; result.rl.Crashes != 1 { + t.Fatalf("list crashes = %d, want 1", result.rl.Crashes) + } + + if restored := debug.SetPanicOnFault(false); restored { + t.Fatal("per-goroutine panic-on-fault mode was not restored") + } +} diff --git a/search/shard_repair.go b/search/shard_repair.go new file mode 100644 index 000000000..454cc3da6 --- /dev/null +++ b/search/shard_repair.go @@ -0,0 +1,135 @@ +package search + +import ( + "errors" + "log" + "runtime/debug" + "sync" + + "github.com/sourcegraph/zoekt" +) + +const shardRepairConcurrency = 4 + +var errShardRepairSuperseded = errors.New("shard repair superseded") + +type shardRepairRequest struct { + key string + faulted zoekt.Searcher +} + +// shardRepairQueue deduplicates repairs by loaded shard instance and bounds +// concurrent open/mmap work when a storage fault affects many shards at once. +type shardRepairQueue struct { + mu sync.Mutex + + entries map[zoekt.Searcher]string + inFlight map[zoekt.Searcher]struct{} + pending []shardRepairRequest + running int + + reload func(string, zoekt.Searcher) error +} + +func newShardRepairQueue(reload func(string, zoekt.Searcher) error) *shardRepairQueue { + return &shardRepairQueue{reload: reload} +} + +func (q *shardRepairQueue) register(searcher zoekt.Searcher, key string) { + if searcher == nil || key == "" { + return + } + + q.mu.Lock() + defer q.mu.Unlock() + if q.entries == nil { + q.entries = make(map[zoekt.Searcher]string) + } + q.entries[searcher] = key +} + +func (q *shardRepairQueue) unregister(searcher zoekt.Searcher) { + if searcher == nil { + return + } + + q.mu.Lock() + defer q.mu.Unlock() + delete(q.entries, searcher) +} + +func (q *shardRepairQueue) schedule(searcher zoekt.Searcher) bool { + q.mu.Lock() + defer q.mu.Unlock() + + key, ok := q.entries[searcher] + if !ok { + // A concurrent watcher replacement or removal already handled this + // shard. + return false + } + if q.inFlight == nil { + q.inFlight = make(map[zoekt.Searcher]struct{}) + } + if _, ok := q.inFlight[searcher]; ok { + return false + } + + q.inFlight[searcher] = struct{}{} + q.pending = append(q.pending, shardRepairRequest{key: key, faulted: searcher}) + q.startPendingLocked() + return true +} + +func (q *shardRepairQueue) startPendingLocked() { + for q.running < shardRepairConcurrency && len(q.pending) > 0 { + request := q.pending[0] + q.pending[0] = shardRepairRequest{} + q.pending = q.pending[1:] + q.running++ + go q.run(request) + } +} + +func (q *shardRepairQueue) run(request shardRepairRequest) { + log.Printf("[WARN] re-opening shard after recovered fault: %s", request.key) + + var ( + reloadErr error + recovered any + stack []byte + ) + func() { + restorePanicOnFault := debug.SetPanicOnFault(true) + defer func() { + debug.SetPanicOnFault(restorePanicOnFault) + if recovered = recover(); recovered != nil { + stack = debug.Stack() + } + }() + reloadErr = q.reload(request.key, request.faulted) + }() + + switch { + case recovered != nil: + log.Printf("[ERROR] fault while re-opening shard %s: %v\n%s", request.key, recovered, stack) + case errors.Is(reloadErr, errShardRepairSuperseded): + log.Printf("[INFO] shard repair superseded by a concurrent update: %s", request.key) + case reloadErr != nil: + log.Printf("[ERROR] failed to re-open shard %s: %v", request.key, reloadErr) + default: + log.Printf("[INFO] re-opened shard after recovered fault: %s", request.key) + } + + q.mu.Lock() + delete(q.inFlight, request.faulted) + q.running-- + q.startPendingLocked() + q.mu.Unlock() +} + +func (q *shardRepairQueue) idle() bool { + q.mu.Lock() + defer q.mu.Unlock() + return len(q.inFlight) == 0 && len(q.pending) == 0 && q.running == 0 +} diff --git a/search/shard_repair_e2e_linux_test.go b/search/shard_repair_e2e_linux_test.go new file mode 100644 index 000000000..321a51918 --- /dev/null +++ b/search/shard_repair_e2e_linux_test.go @@ -0,0 +1,144 @@ +//go:build linux + +package search + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/sourcegraph/zoekt" + "github.com/sourcegraph/zoekt/index" + "github.com/sourcegraph/zoekt/query" +) + +const shardRepairMarker = "UNIQUEMARKERALPHA" + +// TestShardRepairReopensRealShardWithoutWatcher exercises the complete repair +// path against a real shard and filesystem mapping. Replacing the truncated +// shard at a new inode cannot heal the old mapping; only reopening the path can +// make the final search succeed. No DirectoryWatcher runs in this test. +func TestShardRepairReopensRealShardWithoutWatcher(t *testing.T) { + dir := t.TempDir() + shardPath := filepath.Join(dir, "repair_v16.00000.zoekt") + shardBytes := buildShardRepairFixture(t, shardPath) + + ss := newShardedSearcher(1) + (&loader{ss: ss}).load(shardPath) + + q := &query.Substring{Pattern: shardRepairMarker} + search := func() (*zoekt.SearchResult, error) { + return ss.Search(context.Background(), q, &zoekt.SearchOptions{ + ShardMaxMatchCount: 100000, + TotalMaxMatchCount: 100000, + }) + } + + before, err := search() + if err != nil { + t.Fatalf("baseline search: %v", err) + } + if before.Stats.FileCount == 0 { + t.Fatal("baseline search found no fixture documents") + } + + if err := os.Truncate(shardPath, 0); err != nil { + t.Fatalf("truncate shard: %v", err) + } + faulted, err := search() + if err != nil { + t.Fatalf("faulted search returned an error: %v", err) + } + if faulted.Stats.Crashes == 0 { + t.Fatalf( + "truncating the mapped shard induced no fault: FileCount=%d Crashes=%d", + faulted.Stats.FileCount, + faulted.Stats.Crashes, + ) + } + + staged := shardPath + ".staged" + if err := os.WriteFile(staged, shardBytes, 0o644); err != nil { + t.Fatalf("stage replacement shard: %v", err) + } + if err := os.Rename(staged, shardPath); err != nil { + t.Fatalf("publish replacement shard: %v", err) + } + + deadline := time.Now().Add(15 * time.Second) + for { + got, err := search() + if err == nil && got.Stats.Crashes == 0 && got.Stats.FileCount == before.Stats.FileCount { + return + } + if time.Now().After(deadline) { + t.Fatalf( + "shard did not heal: FileCount=%d Crashes=%d, want FileCount=%d Crashes=0", + got.Stats.FileCount, + got.Stats.Crashes, + before.Stats.FileCount, + ) + } + time.Sleep(20 * time.Millisecond) + } +} + +func buildShardRepairFixture(t *testing.T, shardPath string) []byte { + t.Helper() + + buildDir := t.TempDir() + builder, err := index.NewBuilder(index.Options{ + IndexDir: buildDir, + RepositoryDescription: zoekt.Repository{ + Name: "shard-repair-fixture", + }, + DisableCTags: true, + }) + if err != nil { + t.Fatalf("create shard builder: %v", err) + } + + var filler strings.Builder + for i := range 200 { + fmt.Fprintf( + &filler, + "func pad%d() { println(\"payload %d abcdefghij klmnopqrst uvwxyz\") }\n", + i, + i, + ) + } + for i := range 700 { + content := fmt.Sprintf( + "package main\n// %s in file %d\n%s", + shardRepairMarker, + i, + filler.String(), + ) + if err := builder.AddFile(fmt.Sprintf("src/file%d.go", i), []byte(content)); err != nil { + t.Fatalf("add fixture file: %v", err) + } + } + if err := builder.Finish(); err != nil { + t.Fatalf("finish shard: %v", err) + } + + shards, err := filepath.Glob(filepath.Join(buildDir, "*.zoekt")) + if err != nil { + t.Fatalf("find built shard: %v", err) + } + if len(shards) != 1 { + t.Fatalf("built shard count = %d, want 1", len(shards)) + } + data, err := os.ReadFile(shards[0]) + if err != nil { + t.Fatalf("read built shard: %v", err) + } + if err := os.WriteFile(shardPath, data, 0o644); err != nil { + t.Fatalf("write fixture shard: %v", err) + } + return data +} diff --git a/search/shard_repair_test.go b/search/shard_repair_test.go new file mode 100644 index 000000000..4730a7a7d --- /dev/null +++ b/search/shard_repair_test.go @@ -0,0 +1,434 @@ +package search + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/sourcegraph/zoekt" + "github.com/sourcegraph/zoekt/query" +) + +type repairTestSearcher struct { + name string + panicSearch atomic.Bool + panicList atomic.Bool + panicCorrupt atomic.Bool + panicNil atomic.Bool + closed atomic.Bool +} + +type testMemoryFault struct{} + +func (testMemoryFault) Error() string { return "mapped shard fault" } +func (testMemoryFault) Addr() uintptr { return 0x1234 } +func (testMemoryFault) RuntimeError() {} + +func (s *repairTestSearcher) Search(context.Context, query.Q, *zoekt.SearchOptions) (*zoekt.SearchResult, error) { + if s.panicSearch.Load() { + if s.panicNil.Load() { + var pointer *byte + return &zoekt.SearchResult{Stats: zoekt.Stats{Crashes: int(*pointer)}}, nil + } + if s.panicCorrupt.Load() { + panic(errors.New("corrupt shard")) + } + panic(testMemoryFault{}) + } + return &zoekt.SearchResult{}, nil +} + +func (s *repairTestSearcher) List(context.Context, query.Q, *zoekt.ListOptions) (*zoekt.RepoList, error) { + if s.panicList.Load() { + if s.panicCorrupt.Load() { + panic(errors.New("corrupt shard")) + } + panic(testMemoryFault{}) + } + return &zoekt.RepoList{}, nil +} + +func (s *repairTestSearcher) Stats() (*zoekt.RepoStats, error) { + return &zoekt.RepoStats{}, nil +} + +func (s *repairTestSearcher) Close() { s.closed.Store(true) } + +func (s *repairTestSearcher) String() string { return s.name } + +type recordingRepairer struct { + mu sync.Mutex + counts map[string]int + err error + onLoad func(string) +} + +func (r *recordingRepairer) reload(key string, _ zoekt.Searcher) error { + r.mu.Lock() + if r.counts == nil { + r.counts = make(map[string]int) + } + r.counts[key]++ + onLoad, err := r.onLoad, r.err + r.mu.Unlock() + + if onLoad != nil { + onLoad(key) + } + return err +} + +func (r *recordingRepairer) count(key string) int { + r.mu.Lock() + defer r.mu.Unlock() + return r.counts[key] +} + +func waitForRepair(t *testing.T, what string, condition func() bool) { + t.Helper() + deadline := time.Now().Add(10 * time.Second) + for !condition() { + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for %s", what) + } + time.Sleep(2 * time.Millisecond) + } +} + +func testSearcherWithRepairer(repairer *recordingRepairer) (*shardedSearcher, *repairTestSearcher, *rankedShard) { + ss := newShardedSearcher(1) + ss.shardRepairs = newShardRepairQueue(repairer.reload) + searcher := &repairTestSearcher{name: "faulted.zoekt"} + ss.replace(map[string]zoekt.Searcher{searcher.name: searcher}) + return ss, searcher, ss.getLoaded().shards[0] +} + +func TestRecoveredSearchSchedulesShardRepairAndStaysIncomplete(t *testing.T) { + repairer := &recordingRepairer{} + ss, searcher, shard := testSearcherWithRepairer(repairer) + searcher.panicSearch.Store(true) + + result, err := ss.searchOneShard( + context.Background(), + shard, + &query.Const{Value: true}, + &zoekt.SearchOptions{}, + ) + if err != nil { + t.Fatalf("search returned an error: %v", err) + } + if result.Stats.Crashes != 1 { + t.Fatalf("crashes = %d, want 1", result.Stats.Crashes) + } + + waitForRepair(t, "search-triggered repair", ss.shardRepairs.idle) + if got := repairer.count(searcher.name); got != 1 { + t.Fatalf("repair count = %d, want 1", got) + } +} + +func TestRecoveredListSchedulesShardRepair(t *testing.T) { + repairer := &recordingRepairer{} + ss, searcher, shard := testSearcherWithRepairer(repairer) + searcher.panicList.Store(true) + results := make(chan shardListResult, 1) + + ss.listOneShard(context.Background(), shard, &query.Const{Value: true}, nil, results) + result := <-results + if result.err != nil { + t.Fatalf("list returned an error: %v", result.err) + } + if result.rl.Crashes != 1 { + t.Fatalf("crashes = %d, want 1", result.rl.Crashes) + } + + waitForRepair(t, "list-triggered repair", ss.shardRepairs.idle) + if got := repairer.count(searcher.name); got != 1 { + t.Fatalf("repair count = %d, want 1", got) + } +} + +func TestShardRepairsAreDeduplicatedAndConcurrencyBounded(t *testing.T) { + const shardCount = shardRepairConcurrency * 3 + + release := make(chan struct{}) + started := make(chan struct{}, shardCount) + var running atomic.Int32 + var maximum atomic.Int32 + queue := newShardRepairQueue(func(string, zoekt.Searcher) error { + current := running.Add(1) + for { + observed := maximum.Load() + if current <= observed || maximum.CompareAndSwap(observed, current) { + break + } + } + started <- struct{}{} + <-release + running.Add(-1) + return nil + }) + + searchers := make([]*repairTestSearcher, 0, shardCount) + for i := range shardCount { + searcher := &repairTestSearcher{name: fmt.Sprintf("queued-%02d.zoekt", i)} + searchers = append(searchers, searcher) + queue.register(searcher, searcher.name) + if !queue.schedule(searcher) { + t.Fatalf("repair %d was not scheduled", i) + } + if queue.schedule(searcher) { + t.Fatalf("duplicate repair %d was scheduled", i) + } + } + + for range shardRepairConcurrency { + <-started + } + if got := maximum.Load(); got != shardRepairConcurrency { + t.Fatalf("maximum concurrent repairs = %d, want %d", got, shardRepairConcurrency) + } + if queue.idle() { + t.Fatal("queue reported idle with repairs pending") + } + + close(release) + waitForRepair(t, "queued repairs", queue.idle) + if got := maximum.Load(); got > shardRepairConcurrency { + t.Fatalf("maximum concurrent repairs = %d, exceeds %d", got, shardRepairConcurrency) + } +} + +func TestFailedShardRepairDoesNotWithdrawReadinessAndCanRetry(t *testing.T) { + repairer := &recordingRepairer{err: errors.New("stale file handle")} + ss, _, shard := testSearcherWithRepairer(repairer) + ss.markReady() + + ss.shardRepairs.schedule(shard) + waitForRepair(t, "failed repair", ss.shardRepairs.idle) + if !ss.Ready() { + t.Fatal("failed repair withdrew readiness and prevented traffic-driven retry") + } + + repairer.mu.Lock() + repairer.err = nil + repairer.mu.Unlock() + if !ss.shardRepairs.schedule(shard) { + t.Fatal("failed repair retained its single-flight slot") + } + waitForRepair(t, "successful repair retry", ss.shardRepairs.idle) +} + +func TestShardRepairPanicIsContainedAndCanRetry(t *testing.T) { + panicOnLoad := atomic.Bool{} + panicOnLoad.Store(true) + queue := newShardRepairQueue(func(string, zoekt.Searcher) error { + if panicOnLoad.Load() { + panic(errors.New("mapped shard fault during reload")) + } + return nil + }) + searcher := &repairTestSearcher{name: "reload-fault.zoekt"} + queue.register(searcher, searcher.name) + + queue.schedule(searcher) + waitForRepair(t, "contained repair panic", queue.idle) + + panicOnLoad.Store(false) + if !queue.schedule(searcher) { + t.Fatal("repair panic retained its single-flight slot") + } + waitForRepair(t, "repair retry after panic", queue.idle) +} + +func TestWatcherReplacementCanScheduleSamePathWhileOldRepairRuns(t *testing.T) { + const key = "same-path.zoekt" + oldRelease := make(chan struct{}) + var oldRuns atomic.Int32 + var newRuns atomic.Int32 + old := &repairTestSearcher{name: "old"} + newer := &repairTestSearcher{name: "new"} + queue := newShardRepairQueue(func(_ string, faulted zoekt.Searcher) error { + switch faulted { + case old: + oldRuns.Add(1) + <-oldRelease + case newer: + newRuns.Add(1) + } + return nil + }) + + queue.register(old, key) + if !queue.schedule(old) { + t.Fatal("old generation was not scheduled") + } + waitForRepair(t, "old generation repair start", func() bool { return oldRuns.Load() == 1 }) + + queue.unregister(old) + queue.register(newer, key) + if !queue.schedule(newer) { + t.Fatal("new generation was deduplicated against the old generation") + } + waitForRepair(t, "new generation repair", func() bool { return newRuns.Load() == 1 }) + + close(oldRelease) + waitForRepair(t, "both generation repairs", queue.idle) +} + +func TestShardRepairDoesNotReplaceNewerOrRemovedShard(t *testing.T) { + ss := newShardedSearcher(1) + const key = "race.zoekt" + original := &repairTestSearcher{name: "original"} + newer := &repairTestSearcher{name: "newer"} + replacement := &repairTestSearcher{name: "replacement"} + + ss.replace(map[string]zoekt.Searcher{key: original}) + faulted := ss.getLoaded().shards[0] + ss.replace(map[string]zoekt.Searcher{key: newer}) + if ss.swapIfCurrent(key, faulted, replacement) { + t.Fatal("repair replaced a newer shard") + } + + ss.replace(map[string]zoekt.Searcher{key: nil}) + if ss.swapIfCurrent(key, faulted, replacement) { + t.Fatal("repair resurrected a removed shard") + } +} + +func TestShardRepairReplacesCurrentShard(t *testing.T) { + ss := newShardedSearcher(1) + const key = "current.zoekt" + original := &repairTestSearcher{name: "original"} + replacement := &repairTestSearcher{name: "replacement"} + + ss.replace(map[string]zoekt.Searcher{key: original}) + faulted := ss.getLoaded().shards[0] + if !ss.swapIfCurrent(key, faulted, replacement) { + t.Fatal("repair did not replace the current faulted shard") + } + + ss.mu.Lock() + got := ss.shards[key] + ss.mu.Unlock() + if got == nil || got.Searcher != zoekt.Searcher(replacement) { + t.Fatalf("installed shard = %v, want replacement", got) + } +} + +func TestCorruptShardPanicIsContainedWithoutRepair(t *testing.T) { + repairer := &recordingRepairer{} + ss, searcher, shard := testSearcherWithRepairer(repairer) + searcher.panicCorrupt.Store(true) + searcher.panicSearch.Store(true) + + result, err := ss.searchOneShard( + context.Background(), + shard, + &query.Const{Value: true}, + &zoekt.SearchOptions{}, + ) + if err != nil { + t.Fatalf("search returned an error: %v", err) + } + if result.Stats.Crashes != 1 { + t.Fatalf("crashes = %d, want 1", result.Stats.Crashes) + } + if got := repairer.count(searcher.name); got != 0 { + t.Fatalf("repair count = %d, want 0", got) + } +} + +func TestNilPointerPanicIsContainedWithoutRepair(t *testing.T) { + repairer := &recordingRepairer{} + ss, searcher, shard := testSearcherWithRepairer(repairer) + searcher.panicNil.Store(true) + searcher.panicSearch.Store(true) + + result, err := ss.searchOneShard( + context.Background(), + shard, + &query.Const{Value: true}, + &zoekt.SearchOptions{}, + ) + if err != nil { + t.Fatalf("search returned an error: %v", err) + } + if result.Stats.Crashes != 1 { + t.Fatalf("crashes = %d, want 1", result.Stats.Crashes) + } + if got := repairer.count(searcher.name); got != 0 { + t.Fatalf("repair count = %d, want 0", got) + } +} + +type panicIndexFile struct { + closed atomic.Bool +} + +func (*panicIndexFile) Read(uint32, uint32) ([]byte, error) { + panic(errors.New("fault while reading index")) +} +func (*panicIndexFile) Size() (uint32, error) { return 4096, nil } +func (f *panicIndexFile) Close() { f.closed.Store(true) } +func (*panicIndexFile) Name() string { return "panic.zoekt" } + +func TestNewSearcherPanicClosesIndexFile(t *testing.T) { + indexFile := &panicIndexFile{} + var recovered any + func() { + defer func() { recovered = recover() }() + _, _ = newSearcherFromIndexFile(indexFile.Name(), indexFile) + }() + if recovered == nil { + t.Fatal("newSearcherFromIndexFile did not panic") + } + if !indexFile.closed.Load() { + t.Fatal("index file remained open after NewSearcher panic") + } +} + +func TestReloadInstallationPanicClosesReplacement(t *testing.T) { + ss := newShardedSearcher(1) + const key = "install-panic.zoekt" + original := &repairTestSearcher{name: "original"} + ss.replace(map[string]zoekt.Searcher{key: original}) + faulted := ss.getLoaded().shards[0] + + replacement := &repairTestSearcher{name: "replacement"} + replacement.panicList.Store(true) + var recovered any + func() { + defer func() { recovered = recover() }() + ss.installReloadedShard(key, faulted, replacement) + }() + if recovered == nil { + t.Fatal("installReloadedShard did not panic") + } + if !replacement.closed.Load() { + t.Fatal("replacement remained open after installation panic") + } +} + +func TestSupersededReloadClosesReplacement(t *testing.T) { + ss := newShardedSearcher(1) + const key = "superseded.zoekt" + original := &repairTestSearcher{name: "original"} + newer := &repairTestSearcher{name: "newer"} + replacement := &repairTestSearcher{name: "replacement"} + + ss.replace(map[string]zoekt.Searcher{key: original}) + faulted := ss.getLoaded().shards[0] + ss.replace(map[string]zoekt.Searcher{key: newer}) + + if ss.installReloadedShard(key, faulted, replacement) { + t.Fatal("superseded replacement was installed") + } + if !replacement.closed.Load() { + t.Fatal("superseded replacement remained open") + } +} diff --git a/search/shards.go b/search/shards.go index 962effa4f..5b8eb8c39 100644 --- a/search/shards.go +++ b/search/shards.go @@ -213,8 +213,9 @@ type shardedSearcher struct { mu sync.Mutex // protects writes to shards shards map[string]*rankedShard - ready atomic.Bool - ranked atomic.Value + ready atomic.Bool + ranked atomic.Value + shardRepairs *shardRepairQueue } func newShardedSearcher(n int64) *shardedSearcher { @@ -222,6 +223,7 @@ func newShardedSearcher(n int64) *shardedSearcher { shards: make(map[string]*rankedShard), sched: newScheduler(n), } + ss.shardRepairs = newShardRepairQueue(ss.reloadShardByKey) return ss } @@ -584,7 +586,7 @@ func (ss *shardedSearcher) Search(ctx context.Context, q query.Q, opts *zoekt.Se start = time.Now() loaded := ss.getLoaded() - done, err := streamSearch(ctx, proc, q, opts, loaded.shards, collectSender) + done, err := ss.streamSearch(ctx, proc, q, opts, loaded.shards, collectSender) defer done() if err != nil { return nil, err @@ -672,7 +674,7 @@ func (ss *shardedSearcher) StreamSearch(ctx context.Context, q query.Q, opts *zo sender, flush := newFlushCollectSender(opts, sender) - done, err := streamSearch(ctx, proc, q, opts, shards, sender) + done, err := ss.streamSearch(ctx, proc, q, opts, shards, sender) // Even though streaming is done, we may have results sitting in a buffer we // need to flush. So we need to send those before calling done. @@ -690,7 +692,7 @@ func (ss *shardedSearcher) StreamSearch(ctx context.Context, q query.Q, opts *zo // collector can't see. Calling done informs the garbage collector it is free // to collect those shards. The caller must call copyFiles on any // SearchResults it returns/streams out before calling done. -func streamSearch(ctx context.Context, proc *process, q query.Q, opts *zoekt.SearchOptions, shards []*rankedShard, sender zoekt.Sender) (done func(), err error) { +func (ss *shardedSearcher) streamSearch(ctx context.Context, proc *process, q query.Q, opts *zoekt.SearchOptions, shards []*rankedShard, sender zoekt.Sender) (done func(), err error) { tr, ctx := trace.New(ctx, "shardedSearcher.streamSearch", "") overallStart := time.Now() metricSearchRunning.Inc() @@ -756,7 +758,7 @@ func streamSearch(ctx context.Context, proc *process, q query.Q, opts *zoekt.Sea go func() { defer wg.Done() for s := range search { - sr, err := searchOneShard(ctx, s, q, opts) + sr, err := ss.searchOneShard(ctx, s, q, opts) r := &result{priority: s.priority, SearchResult: sr, err: err} results <- r } @@ -968,9 +970,19 @@ func logShardCrash(operation string, s zoekt.Searcher, q query.Q, recovered any, shardRecoveryLogger().Error("crashed shard", fields...) } -func searchOneShard(ctx context.Context, s zoekt.Searcher, q query.Q, opts *zoekt.SearchOptions) (sr *zoekt.SearchResult, err error) { +func isMemoryFault(recovered any) bool { + _, ok := recovered.(interface { + runtime.Error + Addr() uintptr + }) + return ok +} + +func (ss *shardedSearcher) searchOneShard(ctx context.Context, s zoekt.Searcher, q query.Q, opts *zoekt.SearchOptions) (sr *zoekt.SearchResult, err error) { metricSearchShardRunning.Inc() + restorePanicOnFault := debug.SetPanicOnFault(true) defer func() { + debug.SetPanicOnFault(restorePanicOnFault) metricSearchShardRunning.Dec() if e := recover(); e != nil { logShardCrash("search", s, q, e, debug.Stack()) @@ -979,6 +991,9 @@ func searchOneShard(ctx context.Context, s zoekt.Searcher, q query.Q, opts *zoek sr = &zoekt.SearchResult{} } sr.Stats.Crashes = 1 + if isMemoryFault(e) { + ss.shardRepairs.schedule(s) + } } }() @@ -990,12 +1005,17 @@ type shardListResult struct { err error } -func listOneShard(ctx context.Context, s zoekt.Searcher, q query.Q, opts *zoekt.ListOptions, sink chan shardListResult) { +func (ss *shardedSearcher) listOneShard(ctx context.Context, s zoekt.Searcher, q query.Q, opts *zoekt.ListOptions, sink chan shardListResult) { metricListShardRunning.Inc() + restorePanicOnFault := debug.SetPanicOnFault(true) defer func() { + debug.SetPanicOnFault(restorePanicOnFault) metricListShardRunning.Dec() if r := recover(); r != nil { logShardCrash("list", s, q, r, debug.Stack()) + if isMemoryFault(r) { + ss.shardRepairs.schedule(s) + } sink <- shardListResult{ &zoekt.RepoList{Crashes: 1}, nil, } @@ -1075,7 +1095,7 @@ func (ss *shardedSearcher) List(ctx context.Context, q query.Q, opts *zoekt.List for range runtime.GOMAXPROCS(0) { go func() { for s := range feeder { - listOneShard(ctx, s, q, opts, all) + ss.listOneShard(ctx, s, q, opts, all) } }() } @@ -1215,6 +1235,10 @@ func (s *shardedSearcher) replace(shards map[string]zoekt.Searcher) { s.mu.Lock() defer s.mu.Unlock() + s.replaceLocked(shards) +} + +func (s *shardedSearcher) replaceLocked(shards map[string]zoekt.Searcher) { for key, shard := range shards { var r *rankedShard if shard != nil { @@ -1228,6 +1252,13 @@ func (s *shardedSearcher) replace(shards map[string]zoekt.Searcher) { s.shards[key] = r } + if old != nil { + s.shardRepairs.unregister(old) + } + if r != nil { + s.shardRepairs.register(r, key) + } + if old != nil && old.Searcher != nil { // _ ___ /^^\ /^\ /^^\_ // _ _@)@) \ ,,/ '` ~ `'~~ ', `\. @@ -1279,6 +1310,57 @@ func (s *shardedSearcher) replace(shards map[string]zoekt.Searcher) { metricShardsLoaded.Set(float64(len(ranked))) } +func (s *shardedSearcher) reloadShardByKey(key string, faulted zoekt.Searcher) error { + shard, err := loadShard(key) + if err != nil { + return err + } + + if !s.installReloadedShard(key, faulted, shard) { + return errShardRepairSuperseded + } + return nil +} + +func (s *shardedSearcher) installReloadedShard(key string, faulted, shard zoekt.Searcher) (installed bool) { + defer func() { + if !installed { + shard.Close() + } + }() + + installed = s.swapIfCurrent(key, faulted, shard) + return installed +} + +func (s *shardedSearcher) swapIfCurrent(key string, faulted, replacement zoekt.Searcher) bool { + s.mu.Lock() + defer s.mu.Unlock() + + current, ok := s.shards[key] + if !ok || current != faulted { + return false + } + s.replaceLocked(map[string]zoekt.Searcher{key: replacement}) + return true +} + +func newSearcherFromIndexFile(fn string, iFile index.IndexFile) (searcher zoekt.Searcher, err error) { + ownershipTransferred := false + defer func() { + if !ownershipTransferred { + iFile.Close() + } + }() + + searcher, err = index.NewSearcher(iFile) + if err != nil { + return nil, fmt.Errorf("NewSearcher(%s): %v", fn, err) + } + ownershipTransferred = true + return searcher, nil +} + func loadShard(fn string) (zoekt.Searcher, error) { f, err := os.Open(fn) if err != nil { @@ -1289,13 +1371,7 @@ func loadShard(fn string) (zoekt.Searcher, error) { if err != nil { return nil, err } - s, err := index.NewSearcher(iFile) - if err != nil { - iFile.Close() - return nil, fmt.Errorf("NewSearcher(%s): %v", fn, err) - } - - return s, nil + return newSearcherFromIndexFile(fn, iFile) } // prioritySlice is a trivial implementation of an array that provides three