diff --git a/search/aggregate.go b/search/aggregate.go index bc62cef76..13937e4e1 100644 --- a/search/aggregate.go +++ b/search/aggregate.go @@ -173,9 +173,9 @@ func limitSender(cancel context.CancelFunc, sender zoekt.Sender, truncator index }) } -func copyFileSender(sender zoekt.Sender) zoekt.Sender { +func copyFileSender(sender zoekt.Sender, copyResult func(*zoekt.SearchResult)) zoekt.Sender { return zoekt.SenderFunc(func(result *zoekt.SearchResult) { - copyFiles(result) + copyResult(result) sender.Send(result) }) } diff --git a/search/result_copy_fault_linux_test.go b/search/result_copy_fault_linux_test.go new file mode 100644 index 000000000..08bc1a66b --- /dev/null +++ b/search/result_copy_fault_linux_test.go @@ -0,0 +1,405 @@ +//go:build linux + +package search + +import ( + "context" + "os" + "runtime/debug" + "testing" + + "github.com/sourcegraph/zoekt" + "github.com/sourcegraph/zoekt/query" + "golang.org/x/sys/unix" +) + +func truncatedMappedPage(t *testing.T) []byte { + t.Helper() + + file, err := os.CreateTemp(t.TempDir(), "zoekt-copy-fault-*") + if err != nil { + t.Fatalf("create mapped file: %v", err) + } + pageSize := os.Getpagesize() + if err := file.Truncate(int64(pageSize)); err != nil { + t.Fatalf("size mapped file: %v", err) + } + data, err := unix.Mmap( + int(file.Fd()), + 0, + pageSize, + unix.PROT_READ, + unix.MAP_SHARED, + ) + if err != nil { + t.Fatalf("mmap file: %v", err) + } + t.Cleanup(func() { + if err := unix.Munmap(data); err != nil { + t.Errorf("munmap file: %v", err) + } + }) + name := file.Name() + if err := file.Close(); err != nil { + t.Fatalf("close mapped file: %v", err) + } + if err := os.Truncate(name, 0); err != nil { + t.Fatalf("truncate mapped file: %v", err) + } + return data +} + +func TestCopyFilesContainsMmapFaultAndDropsOnlyUnsafeMatch(t *testing.T) { + before := []byte("before") + after := []byte("after") + faulted := truncatedMappedPage(t) + result := &zoekt.SearchResult{ + Files: []zoekt.FileMatch{ + {RepositoryID: 11, Repository: "healthy/before", Content: before}, + {RepositoryID: 22, Repository: "faulted/repo", Content: faulted}, + {RepositoryID: 22, Repository: "faulted/repo", Content: faulted}, + {RepositoryID: 33, Repository: "healthy/after", Content: after}, + }, + } + + origins := &fileMatchOrigins{} + origins.register(&repairTestSearcher{name: "faulted.zoekt"}, result) + faults := copyFiles(result, origins) + if len(faults) != 1 { + t.Fatalf("copy faults = %d, want 1", len(faults)) + } + if faults[0].recovered == nil { + t.Fatal("copyFiles reported a nil fault") + } + if len(faults[0].stack) == 0 { + t.Fatal("copyFiles reported no stack") + } + if result.Stats.Crashes != 1 { + t.Fatalf("crashes = %d, want 1", result.Stats.Crashes) + } + if len(result.Files) != 2 { + t.Fatalf("kept matches = %d, want 2", len(result.Files)) + } + if result.Files[0].RepositoryID != 11 || result.Files[1].RepositoryID != 33 { + t.Fatalf( + "kept repository IDs = [%d %d], want [11 33]", + result.Files[0].RepositoryID, + result.Files[1].RepositoryID, + ) + } + + before[0] = 'X' + after[0] = 'Y' + if string(result.Files[0].Content) != "before" || string(result.Files[1].Content) != "after" { + t.Fatal("healthy match content still aliases its source slices") + } +} + +func TestCopyFileRestoresPanicOnFault(t *testing.T) { + before := debug.SetPanicOnFault(false) + defer debug.SetPanicOnFault(before) + + file := &zoekt.FileMatch{Content: truncatedMappedPage(t)} + recovered, _ := copyFile(file) + if recovered == nil { + t.Fatal("copy of a truncated mapping did not fault") + } + if restored := debug.SetPanicOnFault(false); restored { + t.Fatal("copyFile leaked panic-on-fault mode onto its caller") + } +} + +type copyFaultWiringSearcher struct { + name string + page []byte + allFaulted bool +} + +func (s *copyFaultWiringSearcher) Search(context.Context, query.Q, *zoekt.SearchOptions) (*zoekt.SearchResult, error) { + before := []byte("before") + after := []byte("after") + if s.allFaulted { + before = s.page + after = s.page + } + return &zoekt.SearchResult{ + Files: []zoekt.FileMatch{ + { + RepositoryID: 22, + Repository: "faulted/repository", + FileName: "healthy-before", + Content: before, + Score: 40, + }, + { + RepositoryID: 22, + Repository: "faulted/repository", + FileName: "faulted", + Content: s.page, + Score: 100, + }, + { + RepositoryID: 22, + Repository: "faulted/repository", + FileName: "also-faulted", + Content: s.page, + Score: 90, + }, + { + RepositoryID: 22, + Repository: "faulted/repository", + FileName: "healthy-after", + Content: after, + Score: 30, + }, + }, + Stats: zoekt.Stats{ + FileCount: 4, + MatchCount: 4, + }, + }, nil +} + +func (s *copyFaultWiringSearcher) List(context.Context, query.Q, *zoekt.ListOptions) (*zoekt.RepoList, error) { + return &zoekt.RepoList{ + Repos: []*zoekt.RepoListEntry{{ + Repository: zoekt.Repository{ + ID: 22, + Name: "faulted/repository", + }, + }}, + }, nil +} + +func (s *copyFaultWiringSearcher) Stats() (*zoekt.RepoStats, error) { + return &zoekt.RepoStats{}, nil +} + +func (s *copyFaultWiringSearcher) Close() {} + +func (s *copyFaultWiringSearcher) String() string { return s.name } + +func newCopyFaultWiringSearcher(t *testing.T) (*shardedSearcher, *recordingRepairer) { + t.Helper() + repairer := &recordingRepairer{} + ss := newShardedSearcher(1) + ss.shardRepairs = newShardRepairQueue(repairer.reload) + searcher := ©FaultWiringSearcher{name: "copy-fault.zoekt", page: truncatedMappedPage(t)} + ss.replace(map[string]zoekt.Searcher{searcher.String(): searcher}) + ss.markReady() + return ss, repairer +} + +func assertContainedCopyFault(t *testing.T, result *zoekt.SearchResult, repairer *recordingRepairer) { + t.Helper() + if result.Stats.Crashes != 1 { + t.Fatalf("crashes = %d, want 1", result.Stats.Crashes) + } + if len(result.Files) != 2 { + t.Fatalf("kept matches = %d, want 2", len(result.Files)) + } + for _, file := range result.Files { + if file.FileName == "faulted" { + t.Fatal("faulted match escaped the final copy boundary") + } + } + waitForRepair(t, "copy-fault repair", func() bool { + return repairer.count("copy-fault.zoekt") == 1 + }) +} + +func TestSearchContainsFinalResultCopyFault(t *testing.T) { + ss, repairer := newCopyFaultWiringSearcher(t) + result, err := ss.Search( + context.Background(), + &query.Const{Value: true}, + &zoekt.SearchOptions{}, + ) + if err != nil { + t.Fatalf("Search: %v", err) + } + assertContainedCopyFault(t, result, repairer) +} + +func TestSearchContainsFinalResultCopyFaultWithDisplayLimit(t *testing.T) { + ss, repairer := newCopyFaultWiringSearcher(t) + result, err := ss.Search( + context.Background(), + &query.Const{Value: true}, + &zoekt.SearchOptions{MaxDocDisplayCount: 3}, + ) + if err != nil { + t.Fatalf("Search: %v", err) + } + if result.Stats.Crashes != 1 { + t.Fatalf("crashes = %d, want 1", result.Stats.Crashes) + } + for _, file := range result.Files { + if file.FileName == "faulted" || file.FileName == "also-faulted" { + t.Fatalf("faulted match %q escaped the limited final copy", file.FileName) + } + } + waitForRepair(t, "limited copy-fault repair", func() bool { + return repairer.count("copy-fault.zoekt") == 1 + }) +} + +func TestStreamSearchContainsFinalResultCopyFault(t *testing.T) { + ss, repairer := newCopyFaultWiringSearcher(t) + var copied *zoekt.SearchResult + sender := zoekt.SenderFunc(func(result *zoekt.SearchResult) { + if len(result.Files) > 0 { + copied = result + } + }) + if err := ss.StreamSearch( + context.Background(), + &query.Const{Value: true}, + &zoekt.SearchOptions{}, + sender, + ); err != nil { + t.Fatalf("StreamSearch: %v", err) + } + if copied == nil { + t.Fatal("StreamSearch sent no file results") + } + assertContainedCopyFault(t, copied, repairer) +} + +func TestCopyFaultRepairsExactShardWhenRepositorySpansShards(t *testing.T) { + repairer := &recordingRepairer{} + ss := newShardedSearcher(2) + ss.shardRepairs = newShardRepairQueue(repairer.reload) + + healthy := ©FaultWiringSearcher{name: "healthy.zoekt", page: []byte("healthy")} + faulted := ©FaultWiringSearcher{ + name: "faulted.zoekt", + page: truncatedMappedPage(t), + allFaulted: true, + } + ss.replace(map[string]zoekt.Searcher{ + healthy.name: healthy, + faulted.name: faulted, + }) + ss.markReady() + + result, err := ss.Search( + context.Background(), + &query.Const{Value: true}, + &zoekt.SearchOptions{}, + ) + if err != nil { + t.Fatalf("Search: %v", err) + } + if result.Stats.Crashes != 1 { + t.Fatalf("crashes = %d, want 1", result.Stats.Crashes) + } + waitForRepair(t, "faulted shard repair", ss.shardRepairs.idle) + if got := repairer.count(faulted.name); got != 1 { + t.Fatalf("faulted shard repair count = %d, want 1", got) + } + if got := repairer.count(healthy.name); got != 0 { + t.Fatalf("healthy shard repair count = %d, want 0", got) + } +} + +type multiRepoCopyFaultSearcher struct { + name string + page []byte +} + +func (s *multiRepoCopyFaultSearcher) Search(context.Context, query.Q, *zoekt.SearchOptions) (*zoekt.SearchResult, error) { + return &zoekt.SearchResult{ + Files: []zoekt.FileMatch{ + {RepositoryID: 1, Repository: "org/one", FileName: "one.go", Content: s.page}, + {RepositoryID: 2, Repository: "org/two", FileName: "two.go", Content: s.page}, + }, + RepoURLs: map[string]string{ + "org/one": "", + "org/two": "", + }, + LineFragments: map[string]string{}, + Stats: zoekt.Stats{ + FileCount: 2, + MatchCount: 2, + }, + }, nil +} + +func (s *multiRepoCopyFaultSearcher) List(context.Context, query.Q, *zoekt.ListOptions) (*zoekt.RepoList, error) { + return &zoekt.RepoList{Repos: []*zoekt.RepoListEntry{ + {Repository: zoekt.Repository{ID: 1, Name: "org/one"}}, + {Repository: zoekt.Repository{ID: 2, Name: "org/two"}}, + }}, nil +} + +func (*multiRepoCopyFaultSearcher) Stats() (*zoekt.RepoStats, error) { + return &zoekt.RepoStats{}, nil +} +func (*multiRepoCopyFaultSearcher) Close() {} +func (s *multiRepoCopyFaultSearcher) String() string { return s.name } + +func TestStreamSearchCountsFaultedShardOnceAcrossRepositoryEvents(t *testing.T) { + repairer := &recordingRepairer{} + ss := newShardedSearcher(1) + ss.shardRepairs = newShardRepairQueue(repairer.reload) + searcher := &multiRepoCopyFaultSearcher{ + name: "multi-repo.zoekt", + page: truncatedMappedPage(t), + } + ss.replace(map[string]zoekt.Searcher{searcher.name: searcher}) + ss.markReady() + + var crashes int + sender := zoekt.SenderFunc(func(result *zoekt.SearchResult) { + crashes += result.Stats.Crashes + }) + if err := ss.StreamSearch( + context.Background(), + &query.Const{Value: true}, + &zoekt.SearchOptions{}, + sender, + ); err != nil { + t.Fatalf("StreamSearch: %v", err) + } + if crashes != 1 { + t.Fatalf("aggregate crashes = %d, want 1", crashes) + } + waitForRepair(t, "multi-repository copy-fault repair", ss.shardRepairs.idle) + if got := repairer.count(searcher.name); got != 1 { + t.Fatalf("repair count = %d, want 1", got) + } +} + +func TestAmbiguousCopyFaultDoesNotRepairEitherShard(t *testing.T) { + repairer := &recordingRepairer{} + ss := newShardedSearcher(2) + ss.shardRepairs = newShardRepairQueue(repairer.reload) + page := truncatedMappedPage(t) + first := ©FaultWiringSearcher{name: "first.zoekt", page: page, allFaulted: true} + second := ©FaultWiringSearcher{name: "second.zoekt", page: page, allFaulted: true} + ss.replace(map[string]zoekt.Searcher{ + first.name: first, + second.name: second, + }) + ss.markReady() + + result, err := ss.Search( + context.Background(), + &query.Const{Value: true}, + &zoekt.SearchOptions{}, + ) + if err != nil { + t.Fatalf("Search: %v", err) + } + if result.Stats.Crashes == 0 { + t.Fatal("ambiguous copy fault was not reported") + } + waitForRepair(t, "ambiguous copy-fault handling", ss.shardRepairs.idle) + if got := repairer.count(first.name); got != 0 { + t.Fatalf("first shard repair count = %d, want 0", got) + } + if got := repairer.count(second.name); got != 0 { + t.Fatalf("second shard repair count = %d, want 0", got) + } +} diff --git a/search/result_copy_fault_test.go b/search/result_copy_fault_test.go new file mode 100644 index 000000000..993290078 --- /dev/null +++ b/search/result_copy_fault_test.go @@ -0,0 +1,62 @@ +package search + +import ( + "testing" + + "github.com/sourcegraph/zoekt" +) + +func TestCopyFilesCopiesHealthyMatches(t *testing.T) { + content := []byte("content") + checksum := []byte("checksum") + line := []byte("line") + chunk := []byte("chunk") + result := &zoekt.SearchResult{Files: []zoekt.FileMatch{{ + Content: content, + Checksum: checksum, + LineMatches: []zoekt.LineMatch{{ + Line: line, + }}, + ChunkMatches: []zoekt.ChunkMatch{{ + Content: chunk, + }}, + }}} + + faults := copyFiles(result, &fileMatchOrigins{}) + if len(faults) != 0 { + t.Fatalf("copy faults = %d, want 0", len(faults)) + } + if result.Stats.Crashes != 0 { + t.Fatalf("crashes = %d, want 0", result.Stats.Crashes) + } + + content[0] = 'X' + checksum[0] = 'X' + line[0] = 'X' + chunk[0] = 'X' + file := result.Files[0] + if string(file.Content) != "content" || + string(file.Checksum) != "checksum" || + string(file.LineMatches[0].Line) != "line" || + string(file.ChunkMatches[0].Content) != "chunk" { + t.Fatal("copied match still aliases an mmap-backed source slice") + } +} + +func TestFileMatchOriginsRejectAmbiguousBackingData(t *testing.T) { + shared := []byte("shared") + first := &repairTestSearcher{name: "first.zoekt"} + second := &repairTestSearcher{name: "second.zoekt"} + origins := &fileMatchOrigins{} + origins.register(first, &zoekt.SearchResult{ + Files: []zoekt.FileMatch{{Content: shared}}, + }) + origins.register(second, &zoekt.SearchResult{ + Files: []zoekt.FileMatch{{Content: shared}}, + }) + + file := &zoekt.FileMatch{Content: shared} + if shard, record := origins.recordFault(fileMatchDataPointer(file), file); shard != nil || !record { + t.Fatalf("ambiguous backing data resolved to %v, record = %v", shard, record) + } +} 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..e3a47fc65 100644 --- a/search/shards.go +++ b/search/shards.go @@ -27,6 +27,7 @@ import ( "strconv" "sync" "time" + "unsafe" "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promauto" @@ -213,8 +214,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 +224,7 @@ func newShardedSearcher(n int64) *shardedSearcher { shards: make(map[string]*rankedShard), sched: newScheduler(n), } + ss.shardRepairs = newShardRepairQueue(ss.reloadShardByKey) return ss } @@ -584,7 +587,8 @@ 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) + origins := &fileMatchOrigins{} + done, err := ss.streamSearch(ctx, proc, q, opts, loaded.shards, collectSender, origins) defer done() if err != nil { return nil, err @@ -598,7 +602,7 @@ func (ss *shardedSearcher) Search(ctx context.Context, q query.Q, opts *zoekt.Se } } - copyFiles(aggregate) + ss.handleFileCopyFaults(q, copyFiles(aggregate, origins)) if !loaded.ready { // We may have missed results due to not being fully loaded. @@ -661,7 +665,10 @@ func (ss *shardedSearcher) StreamSearch(ctx context.Context, q query.Q, opts *zo // 4. copyFileSender (copy) // // For streaming, the wrapping has to happen in the inverted order. - sender = copyFileSender(sender) + origins := &fileMatchOrigins{} + sender = copyFileSender(sender, func(result *zoekt.SearchResult) { + ss.handleFileCopyFaults(q, copyFiles(result, origins)) + }) if truncator, hasLimits := index.NewDisplayTruncator(opts); hasLimits { var cancel context.CancelFunc @@ -672,7 +679,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, origins) // 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 +697,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, origins *fileMatchOrigins) (done func(), err error) { tr, ctx := trace.New(ctx, "shardedSearcher.streamSearch", "") overallStart := time.Now() metricSearchRunning.Inc() @@ -756,7 +763,8 @@ 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) + origins.register(s, sr) r := &result{priority: s.priority, SearchResult: sr, err: err} results <- r } @@ -936,17 +944,201 @@ func copySlice(src *[]byte) { *src = dst } -func copyFiles(sr *zoekt.SearchResult) { +func copyFile(file *zoekt.FileMatch) (recovered any, stack []byte) { + restorePanicOnFault := debug.SetPanicOnFault(true) + defer func() { + debug.SetPanicOnFault(restorePanicOnFault) + if recovered = recover(); recovered != nil { + stack = debug.Stack() + } + }() + + copySlice(&file.Content) + copySlice(&file.Checksum) + for i := range file.LineMatches { + copySlice(&file.LineMatches[i].Line) + copySlice(&file.LineMatches[i].Before) + copySlice(&file.LineMatches[i].After) + } + for i := range file.ChunkMatches { + copySlice(&file.ChunkMatches[i].Content) + } + return nil, nil +} + +type fileMatchOrigins struct { + mu sync.Mutex + ranges []fileMatchOriginRange + faulted map[zoekt.Searcher]struct{} + unknown map[unknownFileCopyFaultKey]struct{} +} + +type fileMatchOriginRange struct { + first uintptr + last uintptr + shard zoekt.Searcher +} + +type unknownFileCopyFaultKey struct { + repositoryID uint32 + repository string + fileName string + version string +} + +// Aggregation copies and reorders FileMatch values but preserves their byte +// slice headers. Each shard result contributes one pointer range, keeping +// provenance proportional to searched shards rather than returned matches. +func fileMatchDataPointer(file *zoekt.FileMatch) uintptr { + data := file.Content + if len(data) == 0 { + data = file.Checksum + } + for i := 0; len(data) == 0 && i < len(file.LineMatches); i++ { + data = file.LineMatches[i].Line + if len(data) == 0 { + data = file.LineMatches[i].Before + } + if len(data) == 0 { + data = file.LineMatches[i].After + } + } + for i := 0; len(data) == 0 && i < len(file.ChunkMatches); i++ { + data = file.ChunkMatches[i].Content + } + if len(data) == 0 { + return 0 + } + return uintptr(unsafe.Pointer(unsafe.SliceData(data))) +} + +func (o *fileMatchOrigins) register(shard zoekt.Searcher, result *zoekt.SearchResult) { + if result == nil { + return + } + + var first, last uintptr + for i := range result.Files { + pointer := fileMatchDataPointer(&result.Files[i]) + if pointer == 0 { + continue + } + if first == 0 || pointer < first { + first = pointer + } + if pointer > last { + last = pointer + } + } + if first == 0 { + return + } + + o.mu.Lock() + o.ranges = append(o.ranges, fileMatchOriginRange{ + first: first, + last: last, + shard: shard, + }) + o.mu.Unlock() +} + +func (o *fileMatchOrigins) recordFault(pointer uintptr, file *zoekt.FileMatch) (zoekt.Searcher, bool) { + o.mu.Lock() + defer o.mu.Unlock() + + var shard zoekt.Searcher + for _, candidate := range o.ranges { + if pointer < candidate.first || pointer > candidate.last { + continue + } + if shard != nil && shard != candidate.shard { + shard = nil + break + } + shard = candidate.shard + } + if shard != nil { + if o.faulted == nil { + o.faulted = make(map[zoekt.Searcher]struct{}) + } + if _, exists := o.faulted[shard]; exists { + return shard, false + } + o.faulted[shard] = struct{}{} + return shard, true + } + + key := unknownFileCopyFaultKey{ + repositoryID: file.RepositoryID, + repository: file.Repository, + fileName: file.FileName, + version: file.Version, + } + if o.unknown == nil { + o.unknown = make(map[unknownFileCopyFaultKey]struct{}) + } + if _, exists := o.unknown[key]; exists { + return nil, false + } + o.unknown[key] = struct{}{} + return nil, true +} + +type fileCopyFault struct { + file zoekt.FileMatch + shard zoekt.Searcher + recovered any + stack []byte +} + +func copyFiles(sr *zoekt.SearchResult, origins *fileMatchOrigins) []fileCopyFault { + kept := sr.Files[:0] + var faults []fileCopyFault for i := range sr.Files { - copySlice(&sr.Files[i].Content) - copySlice(&sr.Files[i].Checksum) - for l := range sr.Files[i].LineMatches { - copySlice(&sr.Files[i].LineMatches[l].Line) - copySlice(&sr.Files[i].LineMatches[l].Before) - copySlice(&sr.Files[i].LineMatches[l].After) + file := &sr.Files[i] + originKey := fileMatchDataPointer(file) + fileRecovered, fileStack := copyFile(file) + if fileRecovered != nil { + shard, record := origins.recordFault(originKey, file) + if !record { + continue + } + faults = append(faults, fileCopyFault{ + file: *file, + shard: shard, + recovered: fileRecovered, + stack: fileStack, + }) + continue + } + kept = append(kept, *file) + } + clear(sr.Files[len(kept):]) + sr.Files = kept + if len(faults) > 0 { + sr.Stats.Crashes += len(faults) + metricSearchCrashesTotal.Add(float64(len(faults))) + } + return faults +} + +func (ss *shardedSearcher) handleFileCopyFaults(q query.Q, faults []fileCopyFault) { + for _, fault := range faults { + if fault.shard == nil { + shardRecoveryLogger().Error( + "crashed result copy for unknown shard", + sglog.Uint32("repository_id", fault.file.RepositoryID), + sglog.String("repository", fault.file.Repository), + sglog.String("file", fault.file.FileName), + sglog.String("stacktrace", string(fault.stack)), + ) + continue } - for c := range sr.Files[i].ChunkMatches { - copySlice(&sr.Files[i].ChunkMatches[c].Content) + + logShardCrash("copy", fault.shard, q, fault.recovered, fault.stack) + if isMemoryFault(fault.recovered) { + ss.shardRepairs.schedule(fault.shard) } } } @@ -968,9 +1160,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 +1181,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 +1195,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 +1285,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 +1425,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 +1442,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 +1500,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 +1561,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