fix(search): reconcile directory watcher state

search/shards.go

@@ -305,22 +305,23 @@ type loader struct {
 	ss *shardedSearcher
 }
 
-func (tl *loader) load(keys ...string) {
+func (tl *loader) load(keys ...string) []string {
 	// This is called with all keys on startup, so once this function has
 	// finished running shardedSearcher will be ready.
 	defer tl.ss.markReady()
 
 	if len(keys) == 0 {
 		// If there's nothing to load, we exit early here, but we want to mark
 		// ourselves as ready.
-		return
+		return nil
 	}
 
 	var (
 		mu           sync.Mutex     // synchronizes writes to the shards map
 		wg           sync.WaitGroup // used to wait for all shards to load
 		sem          = semaphore.NewWeighted(int64(runtime.GOMAXPROCS(0)))
 		loadedShards = make(map[string]zoekt.Searcher)
+		loadedKeys   = make([]string, 0, len(keys))
 	)
 
 	publishLoaded := func() {
@@ -360,13 +361,15 @@ func (tl *loader) load(keys ...string) {
 
 			mu.Lock()
 			loadedShards[key] = shard
+			loadedKeys = append(loadedKeys, key)
 			mu.Unlock()
 		}(key)
 	}
 
 	wg.Wait()
 
 	publishLoaded()
+	return loadedKeys
 }
 
 func (tl *loader) drop(keys ...string) {

search/watcher.go

@@ -32,7 +32,7 @@ import (
 
 type shardLoader interface {
 	// Load a new file.
-	load(filenames ...string)
+	load(filenames ...string) []string
 	drop(filenames ...string)
 }
 
@@ -186,7 +186,6 @@ func (s *DirectoryWatcher) scan() error {
 	for k, current := range files {
 		if previous, ok := s.files[k]; !ok || !previous.equal(current) {
 			toLoad = append(toLoad, k)
-			s.files[k] = current
 		}
 	}
 
@@ -204,7 +203,11 @@ func (s *DirectoryWatcher) scan() error {
 	}
 
 	s.loader.drop(toDrop...)
-	s.loader.load(toLoad...)
+	for _, loaded := range s.loader.load(toLoad...) {
+		if current, ok := files[loaded]; ok {
+			s.files[loaded] = current
+		}
+	}
 
 	return nil
 }
@@ -231,6 +234,12 @@ func (s *DirectoryWatcher) watch() error {
 		return err
 	}
 	if err := watcher.Add(s.dir); err != nil {
+		watcher.Close()
+		return err
+	}
+	// Reconcile changes made between the initial scan and watcher registration.
+	if err := s.scan(); err != nil {
+		watcher.Close()
 		return err
 	}
 
@@ -239,6 +248,9 @@ func (s *DirectoryWatcher) watch() error {
 	signal := make(chan struct{}, 1)
 
 	go func() {
+		defer watcher.Close()
+		defer close(signal)
+
 		notify := func() {
 			select {
 			case signal <- struct{}{}:
@@ -247,10 +259,14 @@ func (s *DirectoryWatcher) watch() error {
 		}
 
 		ticker := time.NewTicker(time.Minute)
+		defer ticker.Stop()
 
 		for {
 			select {
-			case event := <-watcher.Events:
+			case event, ok := <-watcher.Events:
+				if !ok {
+					return
+				}
 				// Only notify if a file we read in has changed. This is important to
 				// avoid all the events writing to temporary files.
 				if strings.HasSuffix(event.Name, ".zoekt") || strings.HasSuffix(event.Name, ".meta") {
@@ -261,17 +277,13 @@ func (s *DirectoryWatcher) watch() error {
 				// Periodically just double check the disk
 				notify()
 
-			case err := <-watcher.Errors:
-				// Ignore ErrEventOverflow since we rely on the presence of events so
-				// safe to ignore.
-				if err != nil && err != fsnotify.ErrEventOverflow {
-					log.Println("[ERROR] watcher error:", err)
+			case err, ok := <-watcher.Errors:
+				if !ok {
+					return
 				}
+				handleWatcherError(err, notify)
 
 			case <-s.quit:
-				watcher.Close()
-				ticker.Stop()
-				close(signal)
 				return
 			}
 		}
@@ -288,3 +300,13 @@ func (s *DirectoryWatcher) watch() error {
 
 	return nil
 }
+
+func handleWatcherError(err error, notify func()) {
+	if err == nil {
+		return
+	}
+	notify()
+	if err != fsnotify.ErrEventOverflow {
+		log.Println("[ERROR] watcher error:", err)
+	}
+}

search/watcher_test.go

@@ -15,12 +15,15 @@
 package search
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
 	"testing"
 	"time"
 
+	"github.com/fsnotify/fsnotify"
+
 	"github.com/sourcegraph/zoekt/index"
 )
 
@@ -29,10 +32,11 @@ type loggingLoader struct {
 	drops chan string
 }
 
-func (l *loggingLoader) load(keys ...string) {
+func (l *loggingLoader) load(keys ...string) []string {
 	for _, key := range keys {
 		l.loads <- key
 	}
+	return keys
 }
 
 func (l *loggingLoader) drop(keys ...string) {
@@ -41,6 +45,24 @@ func (l *loggingLoader) drop(keys ...string) {
 	}
 }
 
+type failingLoader struct {
+	attempts int
+	loads    chan string
+}
+
+func (l *failingLoader) load(keys ...string) []string {
+	l.attempts += len(keys)
+	if l.attempts == 1 {
+		return nil
+	}
+	for _, key := range keys {
+		l.loads <- key
+	}
+	return keys
+}
+
+func (*failingLoader) drop(...string) {}
+
 func advanceFS() {
 	time.Sleep(10 * time.Millisecond)
 }
@@ -101,6 +123,82 @@ func TestDirWatcherReloadsReplacementWithSameModTime(t *testing.T) {
 	}
 }
 
+func TestDirWatcherRetriesFailedLoad(t *testing.T) {
+	dir := t.TempDir()
+	shard := filepath.Join(dir, "foo.zoekt")
+	if err := os.WriteFile(shard, []byte("hello"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	loader := &failingLoader{loads: make(chan string, 1)}
+	dw := &DirectoryWatcher{
+		dir:    dir,
+		files:  map[string]watchedFile{},
+		loader: loader,
+	}
+
+	if err := dw.scan(); err != nil {
+		t.Fatal(err)
+	}
+	if err := dw.scan(); err != nil {
+		t.Fatal(err)
+	}
+	if got := <-loader.loads; got != shard {
+		t.Fatalf("got load %q, want %q", got, shard)
+	}
+
+	if err := dw.scan(); err != nil {
+		t.Fatal(err)
+	}
+	if loader.attempts != 2 {
+		t.Fatalf("got %d load attempts, want 2", loader.attempts)
+	}
+}
+
+func TestDirWatcherReconcilesAfterRegistration(t *testing.T) {
+	dir := t.TempDir()
+	logger := &loggingLoader{
+		loads: make(chan string, 1),
+		drops: make(chan string, 1),
+	}
+	dw := &DirectoryWatcher{
+		dir:     dir,
+		files:   map[string]watchedFile{},
+		loader:  logger,
+		quit:    make(chan struct{}),
+		stopped: make(chan struct{}),
+	}
+
+	if err := dw.scan(); err != nil {
+		t.Fatal(err)
+	}
+	shard := filepath.Join(dir, "foo.zoekt")
+	if err := os.WriteFile(shard, []byte("hello"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	if err := dw.watch(); err != nil {
+		t.Fatal(err)
+	}
+	defer dw.Stop()
+
+	if got := <-logger.loads; got != shard {
+		t.Fatalf("got load %q, want %q", got, shard)
+	}
+}
+
+func TestHandleWatcherErrorReconciles(t *testing.T) {
+	calls := 0
+	notify := func() { calls++ }
+
+	handleWatcherError(nil, notify)
+	handleWatcherError(fsnotify.ErrEventOverflow, notify)
+	handleWatcherError(errors.New("watch failed"), notify)
+
+	if calls != 2 {
+		t.Fatalf("got %d notifications, want 2", calls)
+	}
+}
+
 func writeFileWithModTime(t *testing.T, path, content string, modTime time.Time) {
 	t.Helper()
 	if err := os.WriteFile(path, []byte(content), 0o644); err != nil {