fix(search): repair shards after recovered faults

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")
+	}
+}

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
+}

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
+}