diff --git a/search/branchesrepos_test.go b/search/branchesrepos_test.go index 9b5eea6a4..7b36b060d 100644 --- a/search/branchesrepos_test.go +++ b/search/branchesrepos_test.go @@ -184,32 +184,18 @@ func TestSelectRepoSetBranchesReposBuildsBloom(t *testing.T) { // 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) - } + sel.missProbes = branchesReposMinimumMissProbes - var bloom branchesReposBloom - if !sel.maybeBuildBloom(shards[missShards:], &bloom) { + bloom := sel.maybeBuildBloom(shards[missShards:]) + if bloom == nil { 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) + if !bloom.mayContain(id) { + t.Fatalf("matching shard %d was rejected by filter", id) } } assertBranchesReposUnchanged(t, "bloom filter", q, snapshots) @@ -219,28 +205,6 @@ func TestSelectRepoSetBranchesReposBuildsBloom(t *testing.T) { } } -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") @@ -261,14 +225,9 @@ func TestBranchesReposSelectorSkipsNonProbingShards(t *testing.T) { 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) - } - } + sel.missProbes = branchesReposMinimumMissProbes - var bloom branchesReposBloom - if sel.maybeBuildBloom(shards[missShards:], &bloom) { + if sel.maybeBuildBloom(shards[missShards:]) != nil { t.Fatal("selector built a filter despite only unlisted remaining shards") } if !sel.settled { @@ -283,48 +242,12 @@ func TestBranchesReposSelectorSkipsNonProbingShards(t *testing.T) { 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) { + emptySelector.missProbes = branchesReposMinimumMissProbes + if emptySelector.maybeBuildBloom(shards[missShards:]) != nil { 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) diff --git a/search/shards.go b/search/shards.go index 78be4e4d9..6e718841d 100644 --- a/search/shards.go +++ b/search/shards.go @@ -417,6 +417,10 @@ type branchesReposBloom struct { bits [branchesReposBloomWords]uint64 } +var branchesReposBlooms = sync.Pool{ + New: func() any { return new(branchesReposBloom) }, +} + type branchesReposSelector struct { branches []query.BranchRepos cardinality uint64 @@ -430,7 +434,11 @@ type branchesReposSelector struct { } func newBranchesReposSelector(branches []query.BranchRepos, cardinality uint64) *branchesReposSelector { - return &branchesReposSelector{branches: branches, cardinality: cardinality, preferred: -1} + return &branchesReposSelector{ + branches: branches, + cardinality: cardinality, + settled: !branchesReposCanBuildBloom(cardinality), + } } func branchesReposBloomHash(id uint32) uint32 { @@ -464,66 +472,13 @@ func (b *branchesReposBloom) build(branches []query.BranchRepos) { } } -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 { +func (s *branchesReposSelector) maybeBuildBloom(remaining []*rankedShard) *branchesReposBloom { if s.settled || s.missProbes < branchesReposMinimumMissProbes { - return false + return nil } if len(remaining) < branchesReposMinimumRemainingShards { s.settled = true - return false + return nil } // An unlisted shard skips membership checks during selection. Keep the @@ -532,7 +487,7 @@ func (s *branchesReposSelector) maybeBuildBloom(remaining []*rankedShard, bloom for _, shard := range remaining { if shard.repos == nil { s.settled = true - return false + return nil } if len(shard.repos) > 0 { remainingRepoShards++ @@ -540,7 +495,7 @@ func (s *branchesReposSelector) maybeBuildBloom(remaining []*rankedShard, bloom } if remainingRepoShards < branchesReposMinimumRemainingShards { s.settled = true - return false + return nil } // Building the filter visits each requested repository ID once. Count one @@ -549,12 +504,13 @@ func (s *branchesReposSelector) maybeBuildBloom(remaining []*rankedShard, bloom futureProbes := uint64(remainingRepoShards) * uint64(len(s.branches)) if futureProbes <= s.cardinality { s.settled = true - return false + return nil } + bloom := branchesReposBlooms.Get().(*branchesReposBloom) bloom.build(s.branches) s.settled = true - return true + return bloom } // branchesReposCanBuildBloom reports whether the fixed filter retains at least @@ -563,61 +519,13 @@ 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) { +// selectBranchesRepos allocates the fixed filter only after observed misses +// justify it. +func selectBranchesRepos(shards []*rankedShard, branches []query.BranchRepos, cardinality uint64, setSize int) ([]*rankedShard, bool) { selector := newBranchesReposSelector(branches, cardinality) - var bloom branchesReposBloom - bloomReady := false + var bloom *branchesReposBloom 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, @@ -628,58 +536,73 @@ func selectBranchesReposWithBloom(shards []*rankedShard, branches []query.Branch continue } - var any, all bool - if bloomReady { - any, all = selector.matchesBloom(shard.repos, &bloom) - } else { - any, all = selector.matchesDirect(shard.repos) + any, all := false, true + for _, repo := range shard.repos { + matched := false + if bloom == nil || bloom.mayContain(repo.ID) { + matched = branches[selector.preferred].Repos.Contains(repo.ID) + if !matched { + for i, branch := range branches { + if i == selector.preferred { + continue + } + if branch.Repos.Contains(repo.ID) { + selector.preferred = i + matched = true + break + } + } + if !matched { + selector.missProbes += uint64(len(branches)) + } + } + } + any = any || matched + all = all && matched } - if !bloomReady && !selector.settled && selector.missProbes != lastMissProbes { - bloomReady = selector.maybeBuildBloom(shards[shardIndex+1:], &bloom) - lastMissProbes = selector.missProbes + if !selector.settled && selector.missProbes >= branchesReposMinimumMissProbes && bloom == nil { + bloom = selector.maybeBuildBloom(shards[shardIndex+1:]) } 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 + if bloom != nil { + clear(bloom.bits[:]) + branchesReposBlooms.Put(bloom) } - return false + return filtered, filteredAll } -// 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) +// selectBranchesReposDirect avoids adaptive filter bookkeeping when the fixed +// filter would be saturated. +func selectBranchesReposDirect(shards []*rankedShard, branches []query.BranchRepos, setSize int) ([]*rankedShard, bool) { + preferred := 0 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) + any, all := false, true + for _, repo := range shard.repos { + matched := branches[preferred].Repos.Contains(repo.ID) + if !matched { + for i, branch := range branches { + if i != preferred && branch.Repos.Contains(repo.ID) { + preferred = i + matched = true + break + } + } + } + any = any || matched + all = all && matched + } if any { filtered = append(filtered, shard) filteredAll = filteredAll && all @@ -732,7 +655,6 @@ func doSelectRepoSet(shards []*rankedShard, and *query.And) ([]*rankedShard, que 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) @@ -750,57 +672,8 @@ func doSelectRepoSet(shards []*rankedShard, and *query.And) ([]*rankedShard, que return setQuery.Regexp.MatchString(repo.Name) }) case *query.BranchesRepos: - 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 - } - - // 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 { - 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 - }) + if len(setQuery.List) == 0 { + hasRepos = hasReposForPredicate(func(*zoekt.Repository) bool { return false }) break } @@ -825,40 +698,8 @@ func doSelectRepoSet(shards []*rankedShard, and *query.And) ([]*rankedShard, que 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 - }) - } + branchesForSelector = setQuery.List + branchesCardinality = cardinality case *query.Meta: // Meta queries filter repositories based on metadata fields. // By checking this at the shard level, we can skip entire shards @@ -888,10 +729,10 @@ func doSelectRepoSet(shards []*rankedShard, and *query.And) ([]*rankedShard, que var filtered []*rankedShard var filteredAll bool if branchesForSelector != nil { - if buildBloom { - filtered, filteredAll = selectBranchesReposWithBloom(shards, branchesForSelector, branchesCardinality, setSize) + if branchesReposCanBuildBloom(branchesCardinality) { + filtered, filteredAll = selectBranchesRepos(shards, branchesForSelector, branchesCardinality, setSize) } else { - filtered, filteredAll = selectBranchesReposWithPreferred(shards, branchesForSelector, branchesCardinality, setSize) + filtered, filteredAll = selectBranchesReposDirect(shards, branchesForSelector, setSize) } } else { filtered = make([]*rankedShard, 0, setSize)