From c5a37f5e5f36ce1951893728a2e23b2c9df106f7 Mon Sep 17 00:00:00 2001 From: Andrew Mattie Date: Fri, 14 Aug 2026 10:11:50 -0400 Subject: [PATCH] feat/search: allow custom index file openers Keep mmap as the default while allowing directory searchers to inject an alternative IndexFile implementation for initial loads and watcher reloads. Co-Authored-By: OpenAI Codex --- search/shards.go | 57 +++++++++++++++++++++++++++++++++++++------ search/shards_test.go | 51 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 8 deletions(-) diff --git a/search/shards.go b/search/shards.go index 929331545..880c7194a 100644 --- a/search/shards.go +++ b/search/shards.go @@ -229,10 +229,36 @@ func newShardedSearcher(n int64) *shardedSearcher { return ss } +// IndexFileOpener constructs an index file and takes ownership of f. A +// directory searcher may call the opener concurrently for different shards. +type IndexFileOpener func(f *os.File) (index.IndexFile, error) + +// DirectorySearcherOptions configures how a directory searcher loads shards. +type DirectorySearcherOptions struct { + // IndexFileOpener defaults to index.NewIndexFile. + IndexFileOpener IndexFileOpener +} + +func (o DirectorySearcherOptions) indexFileOpener() IndexFileOpener { + if o.IndexFileOpener != nil { + return o.IndexFileOpener + } + return index.NewIndexFile +} + // NewDirectorySearcher returns a searcher instance that loads all // shards corresponding to a glob into memory. func NewDirectorySearcher(dir string) (zoekt.Streamer, error) { - return newDirectorySearcher(dir, true) + return NewDirectorySearcherWithOptions(dir, DirectorySearcherOptions{}) +} + +// NewDirectorySearcherWithOptions is like NewDirectorySearcher, with control +// over how shard files are opened. +func NewDirectorySearcherWithOptions( + dir string, + opts DirectorySearcherOptions, +) (zoekt.Streamer, error) { + return newDirectorySearcher(dir, true, opts) } // ReadySearcher is a Streamer that reports whether its initial load is @@ -249,13 +275,27 @@ type ReadySearcher interface { // partial availability since that is better than no availability on large // instances. func NewDirectorySearcherFast(dir string) (ReadySearcher, error) { - return newDirectorySearcher(dir, false) + return NewDirectorySearcherFastWithOptions(dir, DirectorySearcherOptions{}) +} + +// NewDirectorySearcherFastWithOptions is like NewDirectorySearcherFast, with +// control over how shard files are opened. +func NewDirectorySearcherFastWithOptions( + dir string, + opts DirectorySearcherOptions, +) (ReadySearcher, error) { + return newDirectorySearcher(dir, false, opts) } -func newDirectorySearcher(dir string, waitUntilReady bool) (ReadySearcher, error) { +func newDirectorySearcher( + dir string, + waitUntilReady bool, + opts DirectorySearcherOptions, +) (ReadySearcher, error) { ss := newShardedSearcher(int64(runtime.GOMAXPROCS(0))) tl := &loader{ - ss: ss, + ss: ss, + indexFileOpener: opts.indexFileOpener(), } dw, err := newDirectoryWatcher(dir, tl) if err != nil { @@ -302,7 +342,8 @@ func (s *directorySearcher) Close() { } type loader struct { - ss *shardedSearcher + ss *shardedSearcher + indexFileOpener IndexFileOpener } func (tl *loader) load(keys ...string) { @@ -350,7 +391,7 @@ func (tl *loader) load(keys ...string) { defer sem.Release(1) defer wg.Done() - shard, err := loadShard(key) + shard, err := loadShard(key, tl.indexFileOpener) if err != nil { metricShardsLoadFailedTotal.Inc() log.Printf("[ERROR] reloading: %s, err %v ", key, err) @@ -1529,13 +1570,13 @@ func (s *shardedSearcher) replace(shards map[string]zoekt.Searcher) { metricShardsLoaded.Set(float64(len(ranked))) } -func loadShard(fn string) (zoekt.Searcher, error) { +func loadShard(fn string, openIndexFile IndexFileOpener) (zoekt.Searcher, error) { f, err := os.Open(fn) if err != nil { return nil, err } - iFile, err := index.NewIndexFile(f) + iFile, err := openIndexFile(f) if err != nil { return nil, err } diff --git a/search/shards_test.go b/search/shards_test.go index a74b10d61..c2a606d7b 100644 --- a/search/shards_test.go +++ b/search/shards_test.go @@ -24,10 +24,12 @@ import ( "log" "math" "os" + "path/filepath" "reflect" "runtime" "sort" "strconv" + "sync/atomic" "testing" "testing/quick" "time" @@ -1413,6 +1415,55 @@ func TestNewDirectorySearcher_empty(t *testing.T) { }) } +func TestDirectorySearcherOptionsIndexFileOpener(t *testing.T) { + dir := t.TempDir() + writeShard := func(name string, repositoryID uint32) { + t.Helper() + builder := testShardBuilder(t, &zoekt.Repository{ + ID: repositoryID, + Name: name, + }, index.Document{ + Name: "file.txt", + Content: []byte("content"), + }) + var contents bytes.Buffer + if err := builder.Write(&contents); err != nil { + t.Fatal(err) + } + if err := os.WriteFile( + filepath.Join(dir, name+".zoekt"), + contents.Bytes(), + 0o600, + ); err != nil { + t.Fatal(err) + } + } + + writeShard("initial", 1) + var opens atomic.Int32 + searcher, err := NewDirectorySearcherWithOptions(dir, DirectorySearcherOptions{ + IndexFileOpener: func(file *os.File) (index.IndexFile, error) { + opens.Add(1) + return index.NewIndexFile(file) + }, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(searcher.Close) + if got := opens.Load(); got != 1 { + t.Fatalf("initial shard opens = %d, want 1", got) + } + + writeShard("added", 2) + deadline := testDeadline(t, 10*time.Second) + if !waitForPredicate(deadline, 10*time.Millisecond, func() bool { + return opens.Load() == 2 + }) { + t.Fatalf("shard opens after watcher reload = %d, want 2", opens.Load()) + } +} + // testDeadline returns the deadline for t, but ensures it is no longer than // maxTimeout away. func testDeadline(t *testing.T, maxTimeout time.Duration) time.Time {