Cache docMatchTree results for RepoIds and Meta

When constructing a complicated match tree, this cache
allows zoekt to not have to recompute conditions that it has
already seen before.

Author: John Mason

index/indexdata.go

@@ -103,8 +103,14 @@ type indexData struct {
 
 	// rawConfigMasks contains the encoded RawConfig for each repository
 	rawConfigMasks []uint8
+
+	// Cache for docMatchTree objects
+	docMatchTreeCache docMatchTreeCache
 }
 
+// docMatchTreeCache is a cache for docMatchTree objects so they don't need to be recomputed
+type docMatchTreeCache map[string]*docMatchTree
+
 type symbolData struct {
 	// symContent stores Symbol.Sym and Symbol.Parent.
 	// TODO we don't need to store Symbol.Sym.

index/matchtree.go

@@ -23,6 +23,8 @@ import (
 	"strings"
 	"unicode/utf8"
 
+	"crypto/sha256"
+
 	"github.com/grafana/regexp"
 
 	"github.com/sourcegraph/zoekt"
@@ -971,6 +973,11 @@ func (d *indexData) newMatchTree(q query.Q, opt matchTreeOpt) (matchTree, error)
 	if q == nil {
 		return nil, fmt.Errorf("got nil (sub)query")
 	}
+
+	if d.docMatchTreeCache == nil {
+		d.docMatchTreeCache = make(docMatchTreeCache)
+	}
+
 	switch s := q.(type) {
 	case *query.Regexp:
 		// RegexpToMatchTreeRecursive tries to distill a matchTree that matches a
@@ -1054,6 +1061,11 @@ func (d *indexData) newMatchTree(q query.Q, opt matchTreeOpt) (matchTree, error)
 		}, nil
 
 	case *query.Meta:
+		cacheKey := queryMetaCacheKey(s.Field, s.Value)
+		if cached, ok := d.docMatchTreeCache[cacheKey]; ok {
+			return cached, nil
+		}
+
 		reposWant := make([]bool, len(d.repoMetaData))
 		for repoIdx, r := range d.repoMetaData {
 			if r.Metadata != nil {
@@ -1063,7 +1075,7 @@ func (d *indexData) newMatchTree(q query.Q, opt matchTreeOpt) (matchTree, error)
 			}
 		}
 
-		return &docMatchTree{
+		mt := &docMatchTree{
 			reason:  "Meta",
 			numDocs: d.numDocs(),
 			predicate: func(docID uint32) bool {
@@ -1073,7 +1085,9 @@ func (d *indexData) newMatchTree(q query.Q, opt matchTreeOpt) (matchTree, error)
 				}
 				return reposWant[repoIdx]
 			},
-		}, nil
+		}
+		d.docMatchTreeCache[cacheKey] = mt
+		return mt, nil
 
 	case *query.Substring:
 		return d.newSubstringMatchTree(s)
@@ -1201,19 +1215,27 @@ func (d *indexData) newMatchTree(q query.Q, opt matchTreeOpt) (matchTree, error)
 		}, nil
 
 	case *query.RepoIDs:
+		cacheKey := queryRepoIdsCacheKey(d.repoMetaData)
+		if cached, ok := d.docMatchTreeCache[cacheKey]; ok {
+			return cached, nil
+		}
+
 		reposWant := make([]bool, len(d.repoMetaData))
 		for repoIdx, r := range d.repoMetaData {
 			if s.Repos.Contains(r.ID) {
 				reposWant[repoIdx] = true
 			}
 		}
-		return &docMatchTree{
+
+		mt := &docMatchTree{
 			reason:  "RepoIDs",
 			numDocs: d.numDocs(),
 			predicate: func(docID uint32) bool {
 				return reposWant[d.repos[docID]]
 			},
-		}, nil
+		}
+		d.docMatchTreeCache[cacheKey] = mt
+		return mt, nil
 
 	case *query.Repo:
 		reposWant := make([]bool, len(d.repoMetaData))
@@ -1435,3 +1457,18 @@ func isRegexpAll(r *syntax.Regexp) bool {
 		return false
 	}
 }
+
+func queryMetaCacheKey(field string, value *regexp.Regexp) string {
+	sum := sha256.Sum256([]byte(fmt.Sprintf("%s:%s", field, value.String())))
+	return fmt.Sprintf("Meta:%x", sum[:])
+}
+
+func queryRepoIdsCacheKey(repos []zoekt.Repository) string {
+	var b strings.Builder
+	for _, r := range repos {
+		b.WriteString(fmt.Sprint(r.ID))
+		b.WriteByte(',')
+	}
+	sum := sha256.Sum256([]byte(b.String()))
+	return fmt.Sprintf("RepoIDs:%x", sum[:])
+}

index/matchtree_test.go

@@ -372,10 +372,18 @@ func TestRepoIDs(t *testing.T) {
 		fileBranchMasks: []uint64{1, 1, 1, 1, 1, 1},
 		repos:           []uint16{0, 0, 1, 2, 3, 3},
 	}
-	mt, err := d.newMatchTree(&query.RepoIDs{Repos: roaring.BitmapOf(1, 3, 99)}, matchTreeOpt{})
+	q := &query.RepoIDs{Repos: roaring.BitmapOf(1, 3, 99)}
+	mt, err := d.newMatchTree(q, matchTreeOpt{})
 	if err != nil {
 		t.Fatal(err)
 	}
+
+	// Check that the docMatchTree cache is populated correctly
+	key := queryRepoIdsCacheKey(d.repoMetaData)
+	if _, ok := d.docMatchTreeCache[key]; !ok {
+		t.Errorf("expected docMatchTreeCache to be populated for key %q", key)
+	}
+
 	want := []uint32{2, 4, 5}
 	for i := range want {
 		nextDoc := mt.nextDoc()
@@ -447,6 +455,12 @@ func TestMetaQueryMatchTree(t *testing.T) {
 		t.Fatalf("failed to build matchTree: %v", err)
 	}
 
+	// Check that the docMatchTree cache is populated correctly
+	key := queryMetaCacheKey("license", regexp.MustCompile("M.T"))
+	if _, ok := d.docMatchTreeCache[key]; !ok {
+		t.Errorf("expected docMatchTreeCache to be populated for key %q", key)
+	}
+
 	var matched []uint32
 	for {
 		doc := mt.nextDoc()
@@ -462,3 +476,52 @@ func TestMetaQueryMatchTree(t *testing.T) {
 		t.Errorf("meta match failed: got %v, want %v", matched, want)
 	}
 }
+
+func Test_queryMetaCacheKey(t *testing.T) {
+	cases := []struct {
+		field   string
+		pattern string
+		wantKey string
+	}{
+		// Generated via:
+		// echo -n 'metaField:foo.*bar' | sha256sum
+		{"metaField", "foo.*bar", "Meta:afc6e783c05767285e8657c92c6af09bd8c72d4c0cabe36614b0b2ba3b697724"},
+		// echo -n 'metaField:foo.*baz' | sha256sum
+		{"metaField", "foo.*baz", "Meta:7c5d6616ad2a00042e3ecb1d55cd4ef1907c5b3c232011e45a7f7ba7e8143b63"},
+		// echo -n 'otherField:foo.*bar' | sha256sum
+		{"otherField", "foo.*bar", "Meta:5761c1b19ae8b1c34c5933c8ddb4fe696d80918184547ad42e4953b15700f0ef"},
+	}
+	for _, tc := range cases {
+		re := regexp.MustCompile(tc.pattern)
+		key := queryMetaCacheKey(tc.field, re)
+		if key != tc.wantKey {
+			t.Errorf("unexpected key for field=%q pattern=%q: got %q, want %q", tc.field, tc.pattern, key, tc.wantKey)
+		}
+	}
+}
+
+func Test_queryRepoIdsCacheKey(t *testing.T) {
+	cases := []struct {
+		repos   []zoekt.Repository
+		wantKey string
+	}{
+		// Generated via:
+		// echo -n '123,456,' | sha256sum
+		{[]zoekt.Repository{{ID: 123}, {ID: 456}}, "RepoIDs:a160b50b57496a46824c7e22f8c7047dbbec38752fa1b066d3f50d9f33baaddc"},
+		// echo -n '456,123,' | sha256sum
+		{[]zoekt.Repository{{ID: 456}, {ID: 123}}, "RepoIDs:1d899c857ed96d50e2ad5a9f1505a4a988a69375ec142c8bd29b1aaa545facfb"},
+		// echo -n '123,456,789,' | sha256sum
+		{[]zoekt.Repository{{ID: 123}, {ID: 456}, {ID: 789}}, "RepoIDs:d2c687720e021d3c3d3b8ae461451e144148d84deca4d45d40523f8501c72c39"},
+	}
+	for _, tc := range cases {
+		key := queryRepoIdsCacheKey(tc.repos)
+		if key != tc.wantKey {
+			t.Errorf("unexpected key for repos=%v: got %q, want %q", tc.repos, key, tc.wantKey)
+		}
+		// Check determinism
+		key2 := queryRepoIdsCacheKey(tc.repos)
+		if key != key2 {
+			t.Errorf("key not deterministic for repos=%v: %q vs %q", tc.repos, key, key2)
+		}
+	}
+}