diff --git a/search/branchesrepos_test.go b/search/branchesrepos_test.go new file mode 100644 index 000000000..9b5eea6a4 --- /dev/null +++ b/search/branchesrepos_test.go @@ -0,0 +1,587 @@ +// Copyright 2026 Sourcegraph +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package search + +import ( + "fmt" + "math/rand/v2" + "testing" + + "github.com/RoaringBitmap/roaring/v2" + + "github.com/sourcegraph/zoekt" + "github.com/sourcegraph/zoekt/query" +) + +const ( + branchesReposBenchmarkBranches = 128 + branchesReposBenchmarkContainers = 2_048 + branchesReposBenchmarkRepos = 4 +) + +func newBranchesReposQuery(branches int) *query.BranchesRepos { + list := make([]query.BranchRepos, branches) + for i := range list { + list[i] = query.BranchRepos{ + Branch: fmt.Sprintf("branch-%d", i), + Repos: roaring.New(), + } + } + return &query.BranchesRepos{List: list} +} + +func newBranchesReposShards(shardCount, reposPerShard int) ([]*rankedShard, [][]uint32) { + shards := make([]*rankedShard, shardCount) + ids := make([][]uint32, shardCount) + + for shard := range shards { + repos := make([]*zoekt.Repository, reposPerShard) + ids[shard] = make([]uint32, reposPerShard) + for repo := range repos { + id := uint32(shard*reposPerShard + repo + 1) + ids[shard][repo] = id + repos[repo] = &zoekt.Repository{ID: id} + } + shards[shard] = &rankedShard{repos: repos} + } + + return shards, ids +} + +func addBranchesReposIDs(q *query.BranchesRepos, branch int, ids []uint32) { + for _, id := range ids { + q.List[branch].Repos.Add(id) + } +} + +func matchingBranchesReposShards(shards []*rankedShard, q *query.BranchesRepos) []*rankedShard { + var matching []*rankedShard + for _, shard := range shards { + if shard.repos == nil { + matching = append(matching, shard) + continue + } + + for _, repo := range shard.repos { + for _, branch := range q.List { + if branch.Repos.Contains(repo.ID) { + matching = append(matching, shard) + goto nextShard + } + } + } + nextShard: + } + return matching +} + +type branchesReposSnapshot struct { + bitmap *roaring.Bitmap + contents *roaring.Bitmap +} + +func snapshotBranchesRepos(q *query.BranchesRepos) []branchesReposSnapshot { + snapshots := make([]branchesReposSnapshot, len(q.List)) + for i, branch := range q.List { + snapshots[i] = branchesReposSnapshot{ + bitmap: branch.Repos, + contents: branch.Repos.Clone(), + } + } + return snapshots +} + +func assertBranchesReposUnchanged(t *testing.T, name string, q *query.BranchesRepos, snapshots []branchesReposSnapshot) { + t.Helper() + + for i, snapshot := range snapshots { + if q.List[i].Repos != snapshot.bitmap { + t.Fatalf("%s: selectRepoSet replaced bitmap for branch %q", name, q.List[i].Branch) + } + if !q.List[i].Repos.Equals(snapshot.contents) { + t.Fatalf("%s: selectRepoSet mutated bitmap for branch %q", name, q.List[i].Branch) + } + } +} + +func assertSelectRepoSetBranchesRepos(t *testing.T, name string, shards []*rankedShard, q *query.BranchesRepos) query.Q { + t.Helper() + + snapshots := snapshotBranchesRepos(q) + want := matchingBranchesReposShards(shards, q) + got, gotQuery := selectRepoSet(shards, q) + + if len(got) != len(want) { + t.Fatalf("%s: selected %d shards, want %d", name, len(got), len(want)) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("%s: selected shard %d = %p, want %p", name, i, got[i], want[i]) + } + } + assertBranchesReposUnchanged(t, name, q, snapshots) + + return gotQuery +} + +func TestSelectRepoSetBranchesRepos(t *testing.T) { + shards, ids := newBranchesReposShards(5, 2) + shards[4].repos = nil + + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + addBranchesReposIDs(q, 0, ids[0]) + addBranchesReposIDs(q, branchesReposBenchmarkBranches-1, ids[1]) + addBranchesReposIDs(q, 4, []uint32{ids[2][0]}) + addBranchesReposIDs(q, branchesReposBenchmarkBranches-1, []uint32{ids[2][0]}) + + if gotQuery := assertSelectRepoSetBranchesRepos(t, "mixed membership", shards, q); gotQuery != q { + t.Fatalf("selectRepoSet changed multi-branch query: got %s, want %s", gotQuery, q) + } +} + +func TestSelectRepoSetBranchesReposManyBranches(t *testing.T) { + shards, ids := newBranchesReposShards(branchesReposBenchmarkBranches+1, 1) + shards[len(shards)-1].repos = nil + + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + addBranchesReposIDs(q, 0, ids[0]) + addBranchesReposIDs(q, 4, ids[1]) + addBranchesReposIDs(q, branchesReposBenchmarkBranches-1, ids[2]) + addBranchesReposIDs(q, branchesReposBenchmarkBranches-1, ids[0]) + + if gotQuery := assertSelectRepoSetBranchesRepos(t, "many branches", shards, q); gotQuery != q { + t.Fatalf("selectRepoSet changed multi-branch query: got %s, want %s", gotQuery, q) + } +} + +func TestSelectRepoSetBranchesReposBuildsBloom(t *testing.T) { + // 128 one-repository misses through 128 branch bitmaps make exactly 16,384 + // direct miss probes before the final 16 shards. + const missShards = 128 + shards, ids := newBranchesReposShards(missShards+branchesReposMinimumRemainingShards, 1) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + for shard, shardIDs := range ids[missShards:] { + branch := 0 + if shard%2 != 0 { + branch = branchesReposBenchmarkBranches - 1 + } + addBranchesReposIDs(q, branch, shardIDs) + } + + // Keep the multi-branch query after filtering so its identity can be + // checked below. + shards[missShards].repos = append(shards[missShards].repos, &zoekt.Repository{ID: ids[0][0]}) + + if !branchesReposMaySelect(shards, len(q.List), 2) { + t.Fatal("preflight did not recognize a filterable scan") + } + + snapshots := snapshotBranchesRepos(q) + sel := newBranchesReposSelector(q.List, 2) + for _, shard := range shards[:missShards] { + if sel.containsDirect(shard.repos[0].ID) { + t.Fatalf("miss shard %d matched", shard.repos[0].ID) + } + } + if got := sel.missProbes; got != branchesReposMinimumMissProbes { + t.Fatalf("direct miss probes = %d, want %d", got, branchesReposMinimumMissProbes) + } + + var bloom branchesReposBloom + if !sel.maybeBuildBloom(shards[missShards:], &bloom) { + t.Fatal("selector did not build a membership filter") + } + if bloom.mayContain(ids[0][0]) && sel.containsDirect(ids[0][0]) { + t.Fatalf("miss repository %d matched after filter setup", ids[0][0]) + } + for _, shard := range shards[missShards:] { + id := shard.repos[0].ID + if !bloom.mayContain(id) || !sel.containsDirect(id) { + t.Fatalf("matching shard %d did not match after filter setup", id) + } + } + assertBranchesReposUnchanged(t, "bloom filter", q, snapshots) + + if gotQuery := assertSelectRepoSetBranchesRepos(t, "bloom filter", shards, q); gotQuery != q { + t.Fatalf("selectRepoSet changed multi-branch query: got %s, want %s", gotQuery, q) + } +} + +func TestBranchesReposMayReachSelector(t *testing.T) { + shards, _ := newBranchesReposShards(32, 1) + if branchesReposMayReachSelector(shards, branchesReposBenchmarkBranches) { + t.Fatal("short scan reached selector threshold") + } + + shards, _ = newBranchesReposShards(144, 1) + if !branchesReposMayReachSelector(shards, branchesReposBenchmarkBranches) { + t.Fatal("long scan did not reach selector threshold") + } +} + +func TestBranchesReposMaySelectRequiresRemainingWork(t *testing.T) { + shards, _ := newBranchesReposShards(128, branchesReposBenchmarkRepos) + if !branchesReposMayReachSelector(shards, branchesReposBenchmarkBranches) { + t.Fatal("preflight did not reach selector threshold") + } + if branchesReposMaySelect(shards, branchesReposBenchmarkBranches, 128*100) { + t.Fatal("preflight selected a selector with too little remaining miss work") + } +} + +func TestBranchesReposCanBuildBloomRejectsSaturatedFilter(t *testing.T) { + if !branchesReposCanBuildBloom(branchesReposBloomMaxCardinality) { + t.Fatal("filter capacity was rejected") + } + if branchesReposCanBuildBloom(branchesReposBloomMaxCardinality + 1) { + t.Fatal("saturated filter was accepted") + } +} + +func TestBranchesReposSelectorSkipsNonProbingShards(t *testing.T) { + const missShards = 128 + shards, _ := newBranchesReposShards(missShards+128, 1) + for _, shard := range shards[missShards:] { + shard.repos = nil + } + + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + q.List[0].Repos.Add(100_000) + q.List[branchesReposBenchmarkBranches-1].Repos.Add(200_000) + sel := newBranchesReposSelector(q.List, 2) + for _, shard := range shards[:missShards] { + if sel.containsDirect(shard.repos[0].ID) { + t.Fatalf("miss shard %d matched", shard.repos[0].ID) + } + } + + var bloom branchesReposBloom + if sel.maybeBuildBloom(shards[missShards:], &bloom) { + t.Fatal("selector built a filter despite only unlisted remaining shards") + } + if !sel.settled { + t.Fatal("selector did not settle without known remaining shards") + } + + if gotQuery := assertSelectRepoSetBranchesRepos(t, "unlisted remaining shards", shards, q); gotQuery != q { + t.Fatalf("selectRepoSet changed multi-branch query: got %s, want %s", gotQuery, q) + } + + for _, shard := range shards[missShards:] { + shard.repos = []*zoekt.Repository{} + } + emptySelector := newBranchesReposSelector(q.List, 2) + for _, shard := range shards[:missShards] { + emptySelector.containsDirect(shard.repos[0].ID) + } + var emptyBloom branchesReposBloom + if emptySelector.maybeBuildBloom(shards[missShards:], &emptyBloom) { + t.Fatal("selector built a filter despite only empty remaining shard lists") + } +} + +func TestBranchesReposSelectorRetainsFirstBranch(t *testing.T) { + shards, ids := newBranchesReposShards(128, 1) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + for _, shardIDs := range ids { + addBranchesReposIDs(q, 0, shardIDs) + } + q.List[1].Repos.Add(100_000) + + sel := newBranchesReposSelector(q.List, 129) + for _, shard := range shards { + if !sel.containsDirect(shard.repos[0].ID) { + t.Fatalf("matching shard %d did not match", shard.repos[0].ID) + } + } + if got, want := sel.preferred, -1; got != want { + t.Fatalf("preferred branch = %d, want %d", got, want) + } + if got := sel.missProbes; got != 0 { + t.Fatalf("direct miss probes = %d, want 0", got) + } + if branchesReposFirstMatchIsLater(shards, q.List) { + t.Fatal("first-branch-only query established a later preference") + } + + late := newBranchesReposQuery(branchesReposBenchmarkBranches) + for _, shardIDs := range ids { + addBranchesReposIDs(late, branchesReposBenchmarkBranches-1, shardIDs) + } + if !branchesReposFirstMatchIsLater(shards, late.List) { + t.Fatal("later-branch query did not establish a preference") + } +} + +func TestSelectRepoSetBranchesReposEmptyBitmaps(t *testing.T) { + shards, _ := newBranchesReposShards(branchesReposBenchmarkBranches+100, 1) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + + gotQuery := assertSelectRepoSetBranchesRepos(t, "empty bitmaps", shards, q) + constant, ok := gotQuery.(*query.Const) + if !ok || constant.Value { + t.Fatalf("empty branch repository set returned %s, want FALSE", gotQuery) + } +} + +func TestSelectRepoSetBranchesReposDifferential(t *testing.T) { + random := rand.New(rand.NewPCG(1, 2)) + + for testCase := range 1_500 { + branches := 2 + random.IntN(15) + shardCount := 1 + random.IntN(32) + if testCase%10 == 0 { + branches = branchesReposBenchmarkBranches + shardCount = branchesReposBenchmarkBranches + 1 + random.IntN(branchesReposBenchmarkBranches) + } + + shards, ids := newBranchesReposShards(shardCount, 1+random.IntN(branchesReposBenchmarkRepos)) + q := newBranchesReposQuery(branches) + for _, shardIDs := range ids { + for _, id := range shardIDs { + if random.IntN(4) == 0 { + addBranchesReposIDs(q, random.IntN(branches), []uint32{id}) + } + } + } + for branch := range q.List { + for range random.IntN(3) { + q.List[branch].Repos.Add(uint32(random.Uint64())) + } + } + for shard := range shards { + if random.IntN(16) == 0 { + shards[shard].repos = nil + } + } + + assertSelectRepoSetBranchesRepos(t, fmt.Sprintf("random case %d", testCase), shards, q) + } +} + +func benchmarkSelectRepoSetBranchesRepos(b *testing.B, shards []*rankedShard, q query.Q, wantShards int) { + b.Helper() + + var filtered []*rankedShard + for b.Loop() { + filtered, _ = selectRepoSet(shards, q) + } + + if got := len(filtered); got != wantShards { + b.Fatalf("selected %d shards, want %d", got, wantShards) + } +} + +func BenchmarkSelectRepoSetBranchesRepos(b *testing.B) { + shards, ids := newBranchesReposShards(10_000, branchesReposBenchmarkRepos) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + + for shard, shardIDs := range ids { + if shard%10 == 0 { + addBranchesReposIDs(q, shard%branchesReposBenchmarkBranches, shardIDs) + } + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, 1_000) +} + +func BenchmarkSelectRepoSetBranchesReposLateMatch(b *testing.B) { + shards, ids := newBranchesReposShards(10_000, branchesReposBenchmarkRepos) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + + for _, shardIDs := range ids { + addBranchesReposIDs(q, branchesReposBenchmarkBranches-1, shardIDs) + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, len(shards)) +} + +func BenchmarkSelectRepoSetBranchesReposMatchingPrefix(b *testing.B) { + shards, ids := newBranchesReposShards(10_000, branchesReposBenchmarkRepos) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + + for _, shardIDs := range ids[:4] { + addBranchesReposIDs(q, branchesReposBenchmarkBranches-1, shardIDs) + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, 4) +} + +func BenchmarkSelectRepoSetBranchesReposOneShard(b *testing.B) { + shards, ids := newBranchesReposShards(1, branchesReposBenchmarkRepos) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + addBranchesReposIDs(q, branchesReposBenchmarkBranches-1, ids[0]) + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, 1) +} + +func BenchmarkSelectRepoSetBranchesReposOverlapping(b *testing.B) { + shards, ids := newBranchesReposShards(32, 1) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + + for branch := range q.List { + for _, shardIDs := range ids { + addBranchesReposIDs(q, branch, shardIDs) + } + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, len(shards)) +} + +func BenchmarkSelectRepoSetBranchesReposOverlappingSingleRepo(b *testing.B) { + shards := make([]*rankedShard, 32) + for i := range shards { + shards[i] = &rankedShard{repos: []*zoekt.Repository{{ID: 1}}} + } + + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + for branch := range q.List { + q.List[branch].Repos.Add(1) + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, len(shards)) +} + +func BenchmarkSelectRepoSetBranchesReposLargeBitmapsFewShards(b *testing.B) { + shards, _ := newBranchesReposShards(10, branchesReposBenchmarkRepos) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + + for branch := range q.List { + for id := uint32(0); id < 100; id++ { + q.List[branch].Repos.Add(100_000 + uint32(branch)*100 + id) + } + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, 0) +} + +// BenchmarkSelectRepoSetBranchesReposLargeBitmapsModerateShards covers the +// adaptive boundary where 128 100-ID bitmaps still have too little remaining +// miss work to repay building a membership filter. +func BenchmarkSelectRepoSetBranchesReposLargeBitmapsModerateShards(b *testing.B) { + shards, _ := newBranchesReposShards(128, branchesReposBenchmarkRepos) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + + for branch := range q.List { + for id := uint32(0); id < 100; id++ { + q.List[branch].Repos.Add(100_000 + uint32(branch)*100 + id) + } + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, 0) +} + +// BenchmarkSelectRepoSetBranchesReposDistributedBitmapsModerateShards covers +// sparse bitmaps with 2,048 roaring containers each. The scan reaches the +// adaptive threshold, but visiting every requested ID cannot repay itself +// across only 128 miss shards. +func BenchmarkSelectRepoSetBranchesReposDistributedBitmapsModerateShards(b *testing.B) { + shards, _ := newBranchesReposShards(128, branchesReposBenchmarkRepos) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + + for branch := range q.List { + for container := uint32(1); container <= branchesReposBenchmarkContainers; container++ { + q.List[branch].Repos.Add(container<<16 | uint32(branch)) + } + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, 0) +} + +// BenchmarkSelectRepoSetBranchesReposUnlistedShards covers a failed shard-list +// lookup after enough known misses to reach the filter threshold. Unlisted +// shards skip membership checks and must not repay filter setup. +func BenchmarkSelectRepoSetBranchesReposUnlistedShards(b *testing.B) { + const knownMissShards = 32 + shards, _ := newBranchesReposShards(10_000, branchesReposBenchmarkRepos) + for _, shard := range shards[knownMissShards:] { + shard.repos = nil + } + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + for branch := range q.List { + for id := uint32(0); id < 100; id++ { + q.List[branch].Repos.Add(100_000 + uint32(branch)*100 + id) + } + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, len(shards)-knownMissShards) +} + +// BenchmarkSelectRepoSetBranchesReposFirstBranchWithDistributedBitmap covers +// a cheap first-branch hit path plus a large nonmatching bitmap. Setup must +// retain the direct path rather than build a saturated filter for every query ID. +func BenchmarkSelectRepoSetBranchesReposFirstBranchWithDistributedBitmap(b *testing.B) { + shards, ids := newBranchesReposShards(10_000, branchesReposBenchmarkRepos) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + for _, shardIDs := range ids { + addBranchesReposIDs(q, 0, shardIDs) + } + for container := uint32(1); container <= branchesReposBenchmarkContainers; container++ { + q.List[1].Repos.Add(container<<16 | 1) + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, len(shards)) +} + +// BenchmarkSelectRepoSetBranchesReposLatePrefixThenFirst covers a query whose +// early final-branch matches are followed by a much larger first-branch run. +// The selector must retain the recently matching branch across that suffix. +func BenchmarkSelectRepoSetBranchesReposLatePrefixThenFirst(b *testing.B) { + shards, ids := newBranchesReposShards(10_000, branchesReposBenchmarkRepos) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + + for _, shardIDs := range ids[:2_000] { + addBranchesReposIDs(q, branchesReposBenchmarkBranches-1, shardIDs) + } + for _, shardIDs := range ids[2_000:] { + addBranchesReposIDs(q, 0, shardIDs) + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, len(shards)) +} + +// BenchmarkSelectRepoSetBranchesReposMissPrefixThenFifth covers an initially +// miss-heavy scan followed by repositories that take the fifth branch's cheap +// direct path. The miss prefix must not make that suffix allocate. +func BenchmarkSelectRepoSetBranchesReposMissPrefixThenFifth(b *testing.B) { + shards, ids := newBranchesReposShards(10_000, branchesReposBenchmarkRepos) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + + // Leave the first 1,400 shards outside every branch bitmap. + for _, shardIDs := range ids[1_400:] { + addBranchesReposIDs(q, 4, shardIDs) + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, len(shards)-1_400) +} + +// BenchmarkSelectRepoSetBranchesReposSampledFirstRepoOnly covers a miss-heavy +// query whose interior sample shards match only through their first repository. +// The observed miss work must still justify filtering the remaining shards. +func BenchmarkSelectRepoSetBranchesReposSampledFirstRepoOnly(b *testing.B) { + shards, ids := newBranchesReposShards(10_000, branchesReposBenchmarkRepos) + q := newBranchesReposQuery(branchesReposBenchmarkBranches) + + for shard, shardIDs := range ids { + if shard%10 == 1 { + addBranchesReposIDs(q, branchesReposBenchmarkBranches-1, shardIDs) + } + } + for _, shard := range []int{2_500, 5_000, 7_500} { + addBranchesReposIDs(q, 0, ids[shard][:1]) + } + + benchmarkSelectRepoSetBranchesRepos(b, shards, q, 1_003) +} diff --git a/search/shards.go b/search/shards.go index 437221689..78be4e4d9 100644 --- a/search/shards.go +++ b/search/shards.go @@ -393,6 +393,301 @@ func (ss *shardedSearcher) Close() { ss.replace(shards) } +const ( + // Wait until direct misses have paid for building a membership filter. + branchesReposMinimumMissProbes uint64 = 16 << 10 + + // An adaptive selector near the end of shard selection cannot repay its + // setup cost. + branchesReposMinimumRemainingShards = 16 + + // The fixed 16 KiB filter avoids allocating a union for each query. Two + // hashes keep false positives low for the miss-heavy searches that reach the + // adaptive path. + branchesReposBloomWords = 2048 + branchesReposBloomBits = branchesReposBloomWords * 64 + branchesReposBloomMask = branchesReposBloomBits - 1 + + // Keep at least four filter bits per requested ID, so most direct misses + // can be rejected before scanning the branch bitmaps. + branchesReposBloomMaxCardinality = branchesReposBloomBits / 4 +) + +type branchesReposBloom struct { + bits [branchesReposBloomWords]uint64 +} + +type branchesReposSelector struct { + branches []query.BranchRepos + cardinality uint64 + preferred int + + missProbes uint64 + + // settled is set once filtering has been enabled or ruled out for the + // remaining shards. + settled bool +} + +func newBranchesReposSelector(branches []query.BranchRepos, cardinality uint64) *branchesReposSelector { + return &branchesReposSelector{branches: branches, cardinality: cardinality, preferred: -1} +} + +func branchesReposBloomHash(id uint32) uint32 { + id ^= id >> 16 + id *= 0x7feb352d + id ^= id >> 15 + id *= 0x846ca68b + return id ^ id>>16 +} + +func (b *branchesReposBloom) add(id uint32) { + first := branchesReposBloomHash(id) + second := branchesReposBloomHash(id ^ 0x9e3779b9) + b.bits[(first&branchesReposBloomMask)>>6] |= uint64(1) << (first & 63) + b.bits[(second&branchesReposBloomMask)>>6] |= uint64(1) << (second & 63) +} + +func (b *branchesReposBloom) mayContain(id uint32) bool { + first := branchesReposBloomHash(id) + second := branchesReposBloomHash(id ^ 0x9e3779b9) + return b.bits[(first&branchesReposBloomMask)>>6]&(uint64(1)<<(first&63)) != 0 && + b.bits[(second&branchesReposBloomMask)>>6]&(uint64(1)<<(second&63)) != 0 +} + +func (b *branchesReposBloom) build(branches []query.BranchRepos) { + for _, branch := range branches { + branch.Repos.Iterate(func(id uint32) bool { + b.add(id) + return true + }) + } +} + +func (s *branchesReposSelector) containsDirect(id uint32) bool { + if s.preferred >= 0 { + if s.branches[s.preferred].Repos.Contains(id) { + return true + } + for i, branch := range s.branches { + if i == s.preferred { + continue + } + if branch.Repos.Contains(id) { + s.preferred = i + return true + } + } + } else { + for i, branch := range s.branches { + if branch.Repos.Contains(id) { + // The first branch is already the direct fast path. Remember only a + // later match so first-branch-heavy queries retain that path. + if i != 0 { + s.preferred = i + } + return true + } + } + } + + // A direct miss checks every branch exactly once, regardless of the + // preferred branch's position. + s.missProbes += uint64(len(s.branches)) + return false +} + +func (s *branchesReposSelector) matchesDirect(repos []*zoekt.Repository) (any, all bool) { + all = true + for _, repo := range repos { + matched := s.containsDirect(repo.ID) + any = any || matched + all = all && matched + } + return any, all +} + +func (s *branchesReposSelector) matchesBloom(repos []*zoekt.Repository, bloom *branchesReposBloom) (any, all bool) { + all = true + for _, repo := range repos { + matched := bloom.mayContain(repo.ID) && s.containsDirect(repo.ID) + any = any || matched + all = all && matched + } + return any, all +} + +func (s *branchesReposSelector) maybeBuildBloom(remaining []*rankedShard, bloom *branchesReposBloom) bool { + if s.settled || s.missProbes < branchesReposMinimumMissProbes { + return false + } + if len(remaining) < branchesReposMinimumRemainingShards { + s.settled = true + return false + } + + // An unlisted shard skips membership checks during selection. Keep the + // uncertain remainder on the direct path rather than building a filter. + remainingRepoShards := 0 + for _, shard := range remaining { + if shard.repos == nil { + s.settled = true + return false + } + if len(shard.repos) > 0 { + remainingRepoShards++ + } + } + if remainingRepoShards < branchesReposMinimumRemainingShards { + s.settled = true + return false + } + + // Building the filter visits each requested repository ID once. Count one + // repository per remaining shard, deliberately underestimating compound + // shards so uncertain scans stay on the direct path. + futureProbes := uint64(remainingRepoShards) * uint64(len(s.branches)) + if futureProbes <= s.cardinality { + s.settled = true + return false + } + + bloom.build(s.branches) + s.settled = true + return true +} + +// branchesReposCanBuildBloom reports whether the fixed filter retains at least +// four bits per requested repository ID. +func branchesReposCanBuildBloom(cardinality uint64) bool { + return cardinality <= branchesReposBloomMaxCardinality +} + +// branchesReposMayReachSelector avoids selector metadata work when the known +// prefix cannot reach the miss-work threshold while enough shards remain to use +// a selector. +func branchesReposMayReachSelector(shards []*rankedShard, branches int) bool { + if branches == 0 || len(shards) <= branchesReposMinimumRemainingShards { + return false + } + + var probes uint64 + for _, shard := range shards[:len(shards)-branchesReposMinimumRemainingShards] { + if shard.repos == nil { + continue + } + probes += uint64(len(shard.repos)) * uint64(branches) + if probes >= branchesReposMinimumMissProbes { + return true + } + } + return false +} + +// branchesReposMaySelect avoids selector setup when the known prefix cannot +// reach the miss-work threshold while enough shards remain to use it. +func branchesReposMaySelect(shards []*rankedShard, branches int, cardinality uint64) bool { + if branches == 0 || len(shards) <= branchesReposMinimumRemainingShards { + return false + } + + var probes uint64 + for shardIndex, shard := range shards[:len(shards)-branchesReposMinimumRemainingShards] { + // An unlisted shard remains selected during filtering, but it has no + // known repositories to contribute direct membership probes here. + if shard.repos == nil { + continue + } + probes += uint64(len(shard.repos)) * uint64(branches) + if probes >= branchesReposMinimumMissProbes { + // The adaptive path only counts one repository per remaining shard, + // so use the same conservative bound before allocating its selector. + remainingProbes := uint64(len(shards)-shardIndex-1) * uint64(branches) + return remainingProbes > cardinality + } + } + return false +} + +// selectBranchesReposWithBloom keeps the fixed filter in a cold helper frame: +// ordinary repository selection does not need to grow its stack for it. +func selectBranchesReposWithBloom(shards []*rankedShard, branches []query.BranchRepos, cardinality uint64, setSize int) ([]*rankedShard, bool) { + selector := newBranchesReposSelector(branches, cardinality) + var bloom branchesReposBloom + bloomReady := false + filtered := make([]*rankedShard, 0, setSize) + filteredAll := true + var lastMissProbes uint64 + for shardIndex, shard := range shards { + if shard.repos == nil { + // repos is nil if we failed to List the shard. This shouldn't happen, + // but if it does we don't know what is in it and must search it without + // simplifying the query. + filtered = append(filtered, shard) + filteredAll = false + continue + } + + var any, all bool + if bloomReady { + any, all = selector.matchesBloom(shard.repos, &bloom) + } else { + any, all = selector.matchesDirect(shard.repos) + } + if !bloomReady && !selector.settled && selector.missProbes != lastMissProbes { + bloomReady = selector.maybeBuildBloom(shards[shardIndex+1:], &bloom) + lastMissProbes = selector.missProbes + } + if any { + filtered = append(filtered, shard) + filteredAll = filteredAll && all + } + } + return filtered, filteredAll +} + +// branchesReposFirstMatchIsLater reports whether the first known repository +// matches a branch after the direct first-branch fast path. +func branchesReposFirstMatchIsLater(shards []*rankedShard, branches []query.BranchRepos) bool { + for _, shard := range shards { + if shard.repos == nil || len(shard.repos) == 0 { + continue + } + for branch, candidate := range branches { + if candidate.Repos.Contains(shard.repos[0].ID) { + return branch != 0 + } + } + return false + } + return false +} + +// selectBranchesReposWithPreferred avoids saturating the fixed filter while +// retaining the last matching branch as a cheap direct path. +func selectBranchesReposWithPreferred(shards []*rankedShard, branches []query.BranchRepos, cardinality uint64, setSize int) ([]*rankedShard, bool) { + selector := newBranchesReposSelector(branches, cardinality) + filtered := make([]*rankedShard, 0, setSize) + filteredAll := true + for _, shard := range shards { + if shard.repos == nil { + // repos is nil if we failed to List the shard. This shouldn't happen, + // but if it does we don't know what is in it and must search it without + // simplifying the query. + filtered = append(filtered, shard) + filteredAll = false + continue + } + + any, all := selector.matchesDirect(shard.repos) + if any { + filtered = append(filtered, shard) + filteredAll = filteredAll && all + } + } + return filtered, filteredAll +} + func selectRepoSet(shards []*rankedShard, q query.Q) ([]*rankedShard, query.Q) { and, ok := q.(*query.And) if ok { @@ -435,6 +730,9 @@ func doSelectRepoSet(shards []*rankedShard, and *query.And) ([]*rankedShard, que for i, c := range and.Children { var setSize int var hasRepos func([]*zoekt.Repository) (bool, bool) + var branchesForSelector []query.BranchRepos + var branchesCardinality uint64 + var buildBloom bool switch setQuery := c.(type) { case *query.RepoSet: setSize = len(setQuery.Set) @@ -452,18 +750,115 @@ func doSelectRepoSet(shards []*rankedShard, and *query.And) ([]*rankedShard, que return setQuery.Regexp.MatchString(repo.Name) }) case *query.BranchesRepos: - for _, br := range setQuery.List { - setSize += int(br.Repos.GetCardinality()) + if !branchesReposMayReachSelector(shards, len(setQuery.List)) { + onlyBranch := -1 + for branch, br := range setQuery.List { + branchCardinality := br.Repos.GetCardinality() + setSize += int(branchCardinality) + if branchCardinality == 0 { + continue + } + if onlyBranch < 0 { + onlyBranch = branch + continue + } + onlyBranch = -2 + for _, remaining := range setQuery.List[branch+1:] { + setSize += int(remaining.Repos.GetCardinality()) + } + break + } + if onlyBranch >= 0 { + repos := setQuery.List[onlyBranch].Repos + hasRepos = hasReposForPredicate(func(repo *zoekt.Repository) bool { + return repos.Contains(repo.ID) + }) + } else { + hasRepos = hasReposForPredicate(func(repo *zoekt.Repository) bool { + for _, br := range setQuery.List { + if br.Repos.Contains(repo.ID) { + return true + } + } + return false + }) + } + break } - hasRepos = hasReposForPredicate(func(repo *zoekt.Repository) bool { + // For a saturated first branch, keep the direct path unless the + // first known repository establishes a later useful preference. + if setQuery.List[0].Repos.GetCardinality() > branchesReposBloomMaxCardinality && + !branchesReposFirstMatchIsLater(shards, setQuery.List) { for _, br := range setQuery.List { - if br.Repos.Contains(repo.ID) { - return true + setSize += int(br.Repos.GetCardinality()) + } + hasRepos = hasReposForPredicate(func(repo *zoekt.Repository) bool { + for _, br := range setQuery.List { + if br.Repos.Contains(repo.ID) { + return true + } } + return false + }) + break + } + + var cardinality uint64 + nonEmptyBranches := 0 + onlyBranch := -1 + for branch, br := range setQuery.List { + branchCardinality := br.Repos.GetCardinality() + setSize += int(branchCardinality) + cardinality += branchCardinality + if branchCardinality != 0 { + nonEmptyBranches++ + onlyBranch = branch } - return false - }) + } + + if nonEmptyBranches == 1 { + repos := setQuery.List[onlyBranch].Repos + hasRepos = hasReposForPredicate(func(repo *zoekt.Repository) bool { + return repos.Contains(repo.ID) + }) + break + } + + if !branchesReposCanBuildBloom(cardinality) { + if branchesReposFirstMatchIsLater(shards, setQuery.List) { + branchesForSelector = setQuery.List + branchesCardinality = cardinality + } else { + hasRepos = hasReposForPredicate(func(repo *zoekt.Repository) bool { + for _, br := range setQuery.List { + if br.Repos.Contains(repo.ID) { + return true + } + } + return false + }) + } + break + } + + maySelect := nonEmptyBranches > 1 && + uint64(len(shards))*uint64(len(setQuery.List)) > cardinality && + branchesReposMaySelect(shards, len(setQuery.List), cardinality) + if maySelect { + branchesForSelector = setQuery.List + branchesCardinality = cardinality + buildBloom = true + } else { + hasRepos = hasReposForPredicate(func(repo *zoekt.Repository) bool { + for _, br := range setQuery.List { + if br.Repos.Contains(repo.ID) { + return true + } + } + return false + }) + } case *query.Meta: // Meta queries filter repositories based on metadata fields. // By checking this at the shard level, we can skip entire shards @@ -490,19 +885,28 @@ func doSelectRepoSet(shards []*rankedShard, and *query.And) ([]*rankedShard, que setSize = len(shards) } - filtered := make([]*rankedShard, 0, setSize) - filteredAll := true - - for _, s := range shards { - if s.repos == nil { - // repos is nil if we failed to List the shard. This shouldn't - // happen, but if it does we don't know what is in it and must search - // it without simplifying the query. - filtered = append(filtered, s) - filteredAll = false - } else if any, all := hasRepos(s.repos); any { - filtered = append(filtered, s) - filteredAll = filteredAll && all + var filtered []*rankedShard + var filteredAll bool + if branchesForSelector != nil { + if buildBloom { + filtered, filteredAll = selectBranchesReposWithBloom(shards, branchesForSelector, branchesCardinality, setSize) + } else { + filtered, filteredAll = selectBranchesReposWithPreferred(shards, branchesForSelector, branchesCardinality, setSize) + } + } else { + filtered = make([]*rankedShard, 0, setSize) + filteredAll = true + for _, s := range shards { + if s.repos == nil { + // repos is nil if we failed to List the shard. This shouldn't + // happen, but if it does we don't know what is in it and must search + // it without simplifying the query. + filtered = append(filtered, s) + filteredAll = false + } else if any, all := hasRepos(s.repos); any { + filtered = append(filtered, s) + filteredAll = filteredAll && all + } } }