perf: cache fs traversal across --watch invocations (#117)

When using caching things like tree-sitter queries the slowest thing
remaining is gazelle traversing the fs, essentially doing `ReadDir` and
parsing the BUILDs in each dir.

This patches gazelle to:
1. `aspect_gazelle_runner` gets a reference to the gazelle fs cache,
save it for next time gazelle runs
2. invalidate saved fs reads based on the list of files changed from
watchman
3. populate that gazelle fs cache based on previous invocations, this is
done before gazelle starts to traverse the fs

### Changes are visible to end-users: no

### Test plan

- Covered by existing test cases
- Manual testing; run `configure --watch` and trigger an incremental
invocation, ensure it preloads the fs cache

Author: jbedard

runner/MODULE.bazel

@@ -42,6 +42,7 @@ archive_override(
     patch_strip = 1,
     patches = [
         "//patches:bazelbuild_bazel-gazelle_aspect-cli.patch",
+        "//patches:bazelbuild_bazel-gazelle_aspect-watchman.patch",
         "//pkg/git:gazelle-gitignore.patch",
     ],
     strip_prefix = "bazel-gazelle-2de7b829fef136795c68d9d3b8644a379693cf55",
@@ -68,6 +69,7 @@ go_deps2.module_override(
     patch_strip = 1,
     patches = [
         "//patches:bazelbuild_bazel-gazelle_aspect-cli.patch",
+        "//patches:bazelbuild_bazel-gazelle_aspect-watchman.patch",
         "//pkg/git:gazelle-gitignore.patch",
     ],
     path = "github.com/bazelbuild/bazel-gazelle",

runner/patches/bazelbuild_bazel-gazelle_aspect-watchman.patch

@@ -0,0 +1,17 @@
+diff --git a/walk/walk.go b/walk/walk.go
+index 0e59bcf..b8d774c 100644
+--- a/walk/walk.go
++++ b/walk/walk.go
+@@ -373,6 +373,12 @@ func newWalker(c *config.Config, cexts []config.Configurer, dirs []string, mode
+ 		relsToVisitSeen: make(map[string]struct{}),
+ 	}
+ 
++	// PATCH(caching): additional aspect-gazelle caching of walker
++	if c.Exts["aspect:walkCache:load"] != nil {
++		walkCacheLoad := c.Exts["aspect:walkCache:load"].(func(m interface{}))
++		walkCacheLoad(&w.cache.entryMap)
++	}
++
+ 	// Asynchronously populate the walker cache in the background.
+ 	go w.populateCache()
+ 

runner/pkg/watchman/BUILD.bazel

@@ -26,6 +26,7 @@ go_test(
     ],
     embed = [":watchman"],
     tags = ["manual"],  # TODO: we need watchman in our runners.
+    deps = ["@gazelle//config"],
 )
 
 # At least make sure watch tests can build despite being tagged as manual

runner/pkg/watchman/cache.go

@@ -8,6 +8,7 @@ import (
 	"path"
 	"path/filepath"
 	"runtime"
+	"strings"
 	"sync"
 
 	"github.com/aspect-build/aspect-gazelle/common/cache"
@@ -34,6 +35,10 @@ type watchmanCache struct {
 	old           map[string]map[string]any
 	new           *sync.Map
 	lastClockSpec string
+
+	// A reference to the walk cache used in the last gazelle invocation
+	// to allow restoration in the next invocation.
+	walkCache *sync.Map
 }
 
 var _ cache.Cache = (*watchmanCache)(nil)
@@ -53,10 +58,14 @@ func NewWatchmanCache(c *config.Config) cache.Cache {
 		log.Fatalf("failed to start the watcher: %v", err)
 	}
 
-	return newWatchmanCache(w, diskCachePath)
+	return newWatchmanCache(c, w, diskCachePath)
 }
 
-func newWatchmanCache(w *WatchmanWatcher, diskCachePath string) *watchmanCache {
+// The walk cache of a previous invocation of gazelle.
+// Must be a global var that persists across multiple gazelle invocations.
+var previousWalkCache *sync.Map
+
+func newWatchmanCache(c *config.Config, w *WatchmanWatcher, diskCachePath string) *watchmanCache {
 	wc := &watchmanCache{
 		w:    w,
 		file: diskCachePath,
@@ -65,6 +74,7 @@ func newWatchmanCache(w *WatchmanWatcher, diskCachePath string) *watchmanCache {
 
 		symlinks: &sync.Map{},
 	}
+	wc.populateWalkCache(c)
 	wc.read()
 
 	runtime.SetFinalizer(wc, closeWatchmanCache)
@@ -76,13 +86,53 @@ func closeWatchmanCache(c *watchmanCache) {
 	c.w.Close()
 }
 
+func (c *watchmanCache) populateWalkCache(cfg *config.Config) {
+	// If a walk cache was provided also provide the loader to copy the cached entries
+	// into any fresh walk cache. This must be invoked from a patched gazelle walk.
+	cfg.Exts["aspect:walkCache:load"] = func(m interface{}) {
+		cc := 0
+
+		newWalkCache := m.(*sync.Map)
+		if c.walkCache != nil {
+			c.walkCache.Range(func(key, value any) bool {
+				cc++
+				newWalkCache.Store(key, value)
+				return true
+			})
+		}
+
+		BazelLog.Debugf("Loaded %d walk cache entries into new walk cache\n", cc)
+
+		// Keep a reference to the walk cache for the new gazelle walk invocation
+		// in case of subsequent gazelle invocations.
+		previousWalkCache = newWalkCache
+	}
+}
+
+func invalidateWalkCache(m *sync.Map, staleD string) {
+	if staleD == "." || staleD == "" {
+		m.Clear()
+		return
+	}
+
+	m.Range(func(key, value any) bool {
+		d := key.(string)
+		// Delete the stale directory and any children that may inherit the state
+		if staleD == d || len(d) > len(staleD) && strings.HasPrefix(d, staleD) && d[len(staleD)] == '/' {
+			m.Delete(key)
+		}
+		return true
+	})
+}
+
 func (c *watchmanCache) read() {
 	cacheReader, err := os.Open(c.file)
 	if err != nil {
 		BazelLog.Tracef("Failed to open cache %q: %v", c.file, err)
 		return
 	}
 	defer cacheReader.Close()
+	defer func() { previousWalkCache = nil }()
 
 	var v cacheState
 
@@ -111,14 +161,20 @@ func (c *watchmanCache) read() {
 		return
 	}
 
-	// Discard entries which have changed since the last cache write.
 	for _, p := range cs.Paths {
+		// Discard entries which have changed since the last cache write.
 		delete(v.Entries, p)
+
+		// Discard any walk cache entries for the removed/changed path and its parents.
+		if previousWalkCache != nil {
+			invalidateWalkCache(previousWalkCache, path.Dir(p))
+		}
 	}
 
 	// Persist the still valid entries as the "old" cache state
 	c.old = v.Entries
 	c.lastClockSpec = cs.ClockSpec
+	c.walkCache = previousWalkCache
 
 	// Persist the fact that all persisted paths are not symlinks.
 	// Only new paths with no cache entries will require a stat call.

runner/pkg/watchman/cache_test.go

@@ -3,6 +3,8 @@ package watchman
 import (
 	"os"
 	"testing"
+
+	"github.com/bazelbuild/bazel-gazelle/config"
 )
 
 func TestLoadOrStoreFile(t *testing.T) {
@@ -22,7 +24,7 @@ func TestLoadOrStoreFile(t *testing.T) {
 	defer w.Stop()
 	defer w.Close()
 
-	c := newWatchmanCache(w, getTempFile(t))
+	c := newWatchmanCache(config.New(), w, getTempFile(t))
 
 	computes := 0
 	compute := func(path string, content []byte) (any, error) {
@@ -73,7 +75,7 @@ func TestLoadOrStoreFileSymlink(t *testing.T) {
 	defer w.Stop()
 	defer w.Close()
 
-	c := newWatchmanCache(w, getTempFile(t))
+	c := newWatchmanCache(config.New(), w, getTempFile(t))
 
 	computes := 0
 	compute := func(path string, content []byte) (any, error) {