From 3fb7e2bed94780326aec12951e16b26e8caadf0d Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Fri, 24 Jul 2026 10:52:52 -0700 Subject: [PATCH] feat(speculation): generator contract and bestfirst impl MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add submitqueue/extension/speculation/generator, the candidate-stream composition point (Generator/PathIterator) the standard Speculator pulls from, plus the bestfirst implementation and mocks. bestfirst ranks each path by the product of its dependencies' landing chances (scored through an injected scorer) and generates lazily at two levels. A cross-head heap holds one real, non-terminal candidate per head — never an optimistic upper bound — so popping it yields the globally best-ranked path; each head is itself a lazy stream that materializes its next candidate only once the current one is consumed. Pulling k candidates materializes O(k) paths however large the space behind them, so a head with twelve unresolved dependencies costs no more at the front of the queue than one with two. Within a head, the best path takes the preferred (likelier) bet on every unresolved dependency, and every other path flips some subset of those choices, each flip multiplying the score by that choice's penalty (alternative over preferred). Subsets are walked by a canonical expansion under which every subset has exactly one parent, so each combination is generated exactly once with no visited set, and a child never outranks its parent — which is what lets the heap pop in true descending order. Also corrects the generator README: the leading path is the preferred-bet path, not the all-included one. The two differ whenever a dependency is likelier to fail than to land, so the old wording was wrong for any below-even dependency. The two heaps (across heads, and within a head) order candidates identically, so they share one generic rankedHeap over the standard library's container/heap rather than repeating the heap.Interface boilerplate twice. Test coverage is 22 cases, including the exact emitted sequence across heads, the preferred bet on a below-even dependency, mixed pinned dependencies dropping out of the search, terminal filtering across every terminal status (and deliberately not filtering the non-terminal ones), the materialization counts behind the laziness claim, deterministic tie-breaking, and a randomized cross-check of the whole walk against brute-force enumeration. The generator README gains a worked walkthrough tracing the frontier pop by pop. A head over the depth bound is now recognised from its dependencies' states alone, before anything is scored. Previously every unfinished dependency was scored and only then was the head found to be over the bound. Scoring can be expensive and a head we skip should not pay for it. One visible consequence: a scorer error on a head that was going to be skipped no longer fails Open. --- .../speculation/generator/BUILD.bazel | 9 + .../extension/speculation/generator/README.md | 19 + .../generator/bestfirst/BUILD.bazel | 26 + .../generator/bestfirst/bestfirst.go | 525 +++++++++++ .../generator/bestfirst/bestfirst_test.go | 876 ++++++++++++++++++ .../speculation/generator/generator.go | 45 + .../speculation/generator/mock/BUILD.bazel | 13 + .../generator/mock/generator_mock.go | 98 ++ 8 files changed, 1611 insertions(+) create mode 100644 submitqueue/extension/speculation/generator/BUILD.bazel create mode 100644 submitqueue/extension/speculation/generator/README.md create mode 100644 submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel create mode 100644 submitqueue/extension/speculation/generator/bestfirst/bestfirst.go create mode 100644 submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go create mode 100644 submitqueue/extension/speculation/generator/generator.go create mode 100644 submitqueue/extension/speculation/generator/mock/BUILD.bazel create mode 100644 submitqueue/extension/speculation/generator/mock/generator_mock.go diff --git a/submitqueue/extension/speculation/generator/BUILD.bazel b/submitqueue/extension/speculation/generator/BUILD.bazel new file mode 100644 index 00000000..2e290235 --- /dev/null +++ b/submitqueue/extension/speculation/generator/BUILD.bazel @@ -0,0 +1,9 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "go_default_library", + srcs = ["generator.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/generator", + visibility = ["//visibility:public"], + deps = ["//submitqueue/entity:go_default_library"], +) diff --git a/submitqueue/extension/speculation/generator/README.md b/submitqueue/extension/speculation/generator/README.md new file mode 100644 index 00000000..5a70cc40 --- /dev/null +++ b/submitqueue/extension/speculation/generator/README.md @@ -0,0 +1,19 @@ +# generator + +The `generator` package is a piece the `standard` `Speculator` is built from: a `Generator` produces the queue's candidate paths as one stream, best first, across all heads. It is **not** controller-facing — the speculate controller only knows the `Speculator` contract, and a different `Speculator` need not split its work this way. So there is no `Config` or `Factory` here; a `Generator` is chosen when the `standard` `Speculator` is constructed. + +`Open` starts the stream over the queue's live batches and their path sets and returns a `PathIterator`. The caller pulls one candidate at a time and the generator does only the work that answer needs. Candidates descend in score, never repeat, and never contradict a known fact. A candidate's score means something only for the current run; scores are never stored. + +## `bestfirst` + +A head waiting on unfinished dependencies has one path per combination of guesses — 2ⁿ of them — while callers want the best few. `bestfirst` hands them out in score order without building the rest. + +Dependencies that already finished are not guesses: landed means included, failed or cancelled means excluded. They stay in the path but drop out of the search. Each unfinished one is a two-way guess whose chance of landing comes from an injected `scorer.Scorer`, remembered per dependency so one shared by several heads is scored once; a dependency that is not a live batch is treated as a coin flip. A path's score is its guesses' chances multiplied together. Heads with more unfinished dependencies than the depth bound are skipped until some resolve, and paths that already finished are passed over — so a failed path drops out and the next-best takes its place. + +**The best path is not "bet everything lands."** Each guess takes the *likelier* outcome, so a dependency that will probably fail is excluded. Every other path flips some of those guesses, and each flip costs a known factor. One consequence worth knowing: flipping two cheap guesses can beat flipping one expensive guess, and the ordering handles that correctly. + +**Laziness at two levels.** Across heads, a heap holds each head's current offer — a real, already-built path, never an estimate — so the top of the heap is the best path in the queue. Within a head the same shape repeats: it keeps a small heap, hands out its best, and builds only the one or two paths that come after it. Pulling *k* candidates builds about *k* paths, whatever the size of the set behind them. + +`bestfirst` discards nothing: pull long enough and every path within the depth bound comes out. It does no conflict relaxation (`dropped` bets) — that is the `Speculator`'s call, not the generator's. + +The ordering trick behind all this — how the walk reaches every path exactly once, in score order, without tracking what it has already seen — is documented with a worked example on `expand` in `bestfirst.go`. The behavior is pinned by tests in `bestfirst_test.go`. diff --git a/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel b/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel new file mode 100644 index 00000000..ff2cdd21 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel @@ -0,0 +1,26 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = ["bestfirst.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/generator/bestfirst", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/scorer:go_default_library", + "//submitqueue/extension/speculation/generator:go_default_library", + ], +) + +go_test( + name = "go_default_test", + srcs = ["bestfirst_test.go"], + embed = [":go_default_library"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/scorer:go_default_library", + "//submitqueue/extension/speculation/generator:go_default_library", + "@com_github_stretchr_testify//assert:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + ], +) diff --git a/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go b/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go new file mode 100644 index 00000000..c9981a56 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go @@ -0,0 +1,525 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// 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 bestfirst hands out speculation paths for a queue, best ones first. +// +// A batch we want to build (a "head") often depends on other batches that have +// not finished yet. Each unfinished dependency could go either way, so there is +// more than one sensible thing to build: one path per combination of guesses. +// A path's score is how likely all its guesses are to come true, which is just +// those chances multiplied together. +// +// The catch is the arithmetic. A head with n unfinished dependencies has 2^n +// paths, and the caller normally wants only the best handful. So this package +// never builds them all. It builds one path per head up front, and builds +// another only when the caller takes one. Asking for three paths costs about +// three paths of work, even when the head technically has 4096. +package bestfirst + +import ( + "container/heap" + "context" + "math" + "sort" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/scorer" + "github.com/uber/submitqueue/submitqueue/extension/speculation/generator" +) + +// defaultLandingChance is what we assume for a dependency we cannot score +// because it is not among the live batches: a coin flip. +const defaultLandingChance = 0.5 + +// gen builds best-first iterators. depthBound is the most unfinished +// dependencies a head may have before we skip it, since its path count doubles +// with every one. scorer says how likely a dependency is to land. +type gen struct { + depthBound int + scorer scorer.Scorer +} + +// New returns a best-first generator.Generator. depthBound caps how many +// unfinished dependencies a head may have; heads above it are skipped until +// some of those dependencies resolve. scorer scores each unfinished dependency +// to get its chance of landing. +func New(depthBound int, sc scorer.Scorer) generator.Generator { + return gen{depthBound: depthBound, scorer: sc} +} + +// Open sets up one stream per eligible head and lines up each head's best path, +// ready to hand out. All the scoring happens here, which is why Next never +// calls the scorer and never fails. +func (g gen) Open(ctx context.Context, batches []entity.Batch, pathSets []entity.SpeculationPathSet) (generator.PathIterator, error) { + batchByID := make(map[string]entity.Batch, len(batches)) + for _, b := range batches { + batchByID[b.ID] = b + } + terminalByHead := terminalPathIDs(pathSets) + + // chanceOf remembers each dependency's chance of landing, so a dependency + // that several heads wait on is only scored once. + chances := make(map[string]float64) + chanceOf := func(depID string) (float64, error) { + if c, ok := chances[depID]; ok { + return c, nil + } + dep, ok := batchByID[depID] + if !ok { + // Not a live batch, so there is nothing to score. Remember the coin + // flip too, so we don't look it up again. + chances[depID] = defaultLandingChance + return defaultLandingChance, nil + } + c, err := g.scorer.Score(ctx, dep) + if err != nil { + return 0, err + } + chances[depID] = c + return c, nil + } + + it := &iterator{heap: &rankedHeap[*headStream]{}} + for _, head := range batches { + // Only Speculating heads get paths. Batches in any other state (created, + // merging, cancelling, finished) still tell us how their dependencies + // turned out, but we never propose work on them. + if head.State != entity.BatchStateSpeculating { + continue + } + s, ok, err := g.newHeadStream(head, batchByID, terminalByHead[head.ID], chanceOf) + if err != nil { + return nil, err + } + if !ok { + continue + } + it.streams = append(it.streams, s) + // Ask for the head's first path. A head whose paths have all finished + // already has nothing to offer and quietly stays out of the merge. + it.queueNext(s) + } + return it, nil +} + +// choice is one unfinished dependency that a head has to guess about. +// +// The likelier of the two outcomes is the "preferred" bet and the other is the +// "alternative". Betting the alternative is always the worse guess, and penalty +// says how much worse: multiply a path's score by it and you get the score of +// the same path with this one guess flipped. A penalty near 1 means flipping +// barely costs anything (the dependency was close to a coin flip); a penalty +// near 0 means flipping it is very expensive. +type choice struct { + // dep is the dependency batch being guessed about. + dep string + // slot is where this dependency sits in the head's bets, so a flip writes + // back to the right place and bets stay in queue order. + slot int + // alternativeBet is the bet used when this guess is flipped. + alternativeBet entity.DependencyBetType + // preferred is the chance of the likelier outcome, so always >= 0.5. + preferred float64 + // alternative is the chance of the other outcome, so always <= 0.5. + alternative float64 + // penalty is alternative/preferred, between 0 and 1. + penalty float64 +} + +// ranked is one entry in a rankedHeap: a path, its ID, and whatever the heap's +// owner wants to keep alongside it. +type ranked[T any] struct { + cand entity.CandidatePath + id string // cand.Path.ID(), kept here so comparing two entries needn't rehash + // payload is the owner's own bookkeeping: which guesses a path flips within + // a head, or which head a path came from across heads. + payload T +} + +// rankedHeap keeps paths in best-first order: highest score first, and when two +// scores tie, lowest path ID first so the order never changes between runs. +// Both heaps in this file want exactly that and differ only in what they carry +// along, so they share this one type. Go's container/heap does the real work +// here; the methods below are just the interface it asks for. +type rankedHeap[T any] []ranked[T] + +func (h rankedHeap[T]) Len() int { return len(h) } +func (h rankedHeap[T]) Less(i, j int) bool { + // Greater-than, not less-than: that is what makes this a max-heap, so the + // highest-scoring path ends up on top. + if h[i].cand.RankingScore != h[j].cand.RankingScore { + return h[i].cand.RankingScore > h[j].cand.RankingScore + } + // Same score: fall back to path ID purely so the order is repeatable. + return h[i].id < h[j].id +} +func (h rankedHeap[T]) Swap(i, j int) { h[i], h[j] = h[j], h[i] } +func (h *rankedHeap[T]) Push(x any) { *h = append(*h, x.(ranked[T])) } +func (h *rankedHeap[T]) Pop() any { + // Careful: this does not pick the best entry. container/heap has already + // swapped the best one to the end of the slice before calling us, so all we + // do is chop the last element off and hand it back. + old := *h + n := len(old) + e := old[n-1] + *h = old[:n-1] + return e +} + +// pushRanked adds an entry to h. +func pushRanked[T any](h *rankedHeap[T], e ranked[T]) { + heap.Push(h, e) +} + +// popRanked takes h's best entry out. It panics on an empty heap, so callers +// check Len first. +func popRanked[T any](h *rankedHeap[T]) ranked[T] { + return heap.Pop(h).(ranked[T]) +} + +// flipSet says which of a head's guesses a path flips. +// +// Every path is described by how it differs from the head's best path. The +// numbers are positions in the head's choices list, marking the guesses that go +// the other way; any choice not listed keeps its preferred bet. +// +// flipSet{} the best path — every guess is the likely one +// flipSet{0} the same, except choices[0] guesses the other way +// flipSet{0, 2} the same, except choices[0] and choices[2] both do +// +// A head with n choices has 2^n possible flip sets, one per path. The numbers +// are always kept in increasing order, which is what makes the walk in expand +// reach each set exactly once. +type flipSet []int + +// combination is one path in a head's search, carrying the flipSet that +// describes it. +type combination = ranked[flipSet] + +// front is the best path a head is currently offering, waiting in the +// cross-head heap, carrying the stream it came from. +type front = ranked[*headStream] + +// headStream hands out one head's paths, best first, building them as it goes. +// +// It starts from the head's best path and works outward by flipping guesses. +// The heap holds paths that have been built but not handed out yet; each time +// one is handed out, the one or two paths that come just after it get built. +// That keeps the heap stocked without ever building the whole set. +// +// Say a head waits on three dependencies with chances d0 0.9, d1 0.8, d2 0.6. +// All three are better than even, so the best path includes all three and +// scores 0.9 × 0.8 × 0.6 = 0.432. Sorted cheapest-flip-first the choices come +// out c0=d2, c1=d1, c2=d0. Pulling all eight paths then goes: +// +// handed out still waiting afterwards +// {} .432 {c0}·.288 +// {c0} .288 {c1}·.108 {c0,c1}·.072 +// {c1} .108 {c0,c1}·.072 {c2}·.048 {c1,c2}·.012 +// {c0,c1} .072 {c2}·.048 {c0,c2}·.032 {c1,c2}·.012 {c0,c1,c2}·.008 +// {c2} .048 {c0,c2}·.032 {c1,c2}·.012 {c0,c1,c2}·.008 +// {c0,c2} .032 {c1,c2}·.012 {c0,c1,c2}·.008 +// {c1,c2} .012 {c0,c1,c2}·.008 +// {c0,c1,c2} .008 — +// +// The scores come out in order even though nothing is ever sorted, and the +// waiting list never holds more than four paths for a set of eight. Stop after +// two and only three paths were ever built. +// +// Look at the fourth row: {c0,c1} flips two guesses yet still beats {c2}, +// which flips only one. That is correct — flipping the two cheap guesses costs +// less than flipping the expensive one. +type headStream struct { + // head is the batch these paths are for. + head string + // bets is the best path's bets, in queue order: the known bet where the + // dependency has already finished, the preferred guess everywhere else. + bets []entity.DependencyBet + // choices are the head's unfinished dependencies, cheapest flip first. + choices []choice + // terminal holds the IDs of this head's paths that already finished, so we + // don't offer them again. + terminal map[string]bool + // heap holds paths built but not yet handed out. + heap *rankedHeap[flipSet] + + // generated counts how many paths have been built. Nothing reads it in + // production; it is here so tests can check we really are building lazily. + generated int +} + +// settledBet reports the bet a dependency's state forces, if that state is +// final. An unfinished dependency has no forced bet and has to be guessed at. +func settledBet(state entity.BatchState) (entity.DependencyBetType, bool) { + switch state { + case entity.BatchStateSucceeded: + // Already landed, so including it is a fact, not a guess. + return entity.BetIncluded, true + case entity.BatchStateFailed, entity.BatchStateCancelled: + // Never landing, so leaving it out is a fact too. + return entity.BetExcluded, true + default: + return entity.BetUnknown, false + } +} + +// newHeadStream sets up one head. It sorts the head's dependencies into two +// piles — ones that already finished, whose bet is a known fact, and unfinished +// ones, which become choices to guess about — and then builds the head's best +// path so the stream has something to offer. It returns ok=false when the head +// has more unfinished dependencies than the depth bound allows. +func (g gen) newHeadStream( + head entity.Batch, + batchByID map[string]entity.Batch, + terminal map[string]bool, + chanceOf func(string) (float64, error), +) (*headStream, bool, error) { + // Count the guesses before making any. Reading a state is free while scoring + // may not be, so a head we are going to skip costs nothing. + unfinished := 0 + for _, dep := range head.Dependencies { + if _, settled := settledBet(batchByID[dep].State); !settled { + unfinished++ + } + } + // Too many things to guess about — skip this head until some of them resolve. + if unfinished > g.depthBound { + return nil, false, nil + } + + bets := make([]entity.DependencyBet, len(head.Dependencies)) + choices := make([]choice, 0, unfinished) + for i, dep := range head.Dependencies { + if bet, settled := settledBet(batchByID[dep].State); settled { + bets[i] = entity.DependencyBet{Batch: dep, Bet: bet} + continue + } + c, err := chanceOf(dep) + if err != nil { + return nil, false, err + } + // Guess the likelier outcome. On an exact coin flip, guess it lands. + preferredBet, alternativeBet := entity.BetIncluded, entity.BetExcluded + if c < defaultLandingChance { + preferredBet, alternativeBet = entity.BetExcluded, entity.BetIncluded + } + // These line up with the bets above: c is the chance it lands, so when + // we bet it lands the bigger number is c, and when we bet it doesn't + // the bigger number is 1-c. Either way preferred is the chance of the + // bet we picked. + preferred, alternative := math.Max(c, 1-c), math.Min(c, 1-c) + bets[i] = entity.DependencyBet{Batch: dep, Bet: preferredBet} + choices = append(choices, choice{ + dep: dep, + slot: i, + alternativeBet: alternativeBet, + preferred: preferred, + alternative: alternative, + // preferred is at least 0.5, so this never divides by zero. + penalty: alternative / preferred, + }) + } + // Cheapest flip first — the guess we'd mind changing least. expand depends on + // this order: it only ever moves a flip to the next position along, so having + // them sorted is what guarantees a path never scores above the one it came + // from. + sort.Slice(choices, func(i, j int) bool { + if choices[i].penalty != choices[j].penalty { + return choices[i].penalty > choices[j].penalty + } + return choices[i].slot < choices[j].slot + }) + + s := &headStream{ + head: head.ID, + bets: bets, + choices: choices, + terminal: terminal, + heap: &rankedHeap[flipSet]{}, + } + // Build the head's best path. An empty flip set flips nothing, so every guess + // is the likely one. This is the only path built up front; the rest grow out + // of it as the caller pulls. + s.push(flipSet{}) + return s, true, nil +} + +// Next hands out this head's next-best path. On the way out it builds whatever +// comes after that path, so the heap always has the next candidates ready. It +// returns ok=false once the head has nothing left to offer. +func (s *headStream) Next() (entity.CandidatePath, bool) { + for s.heap.Len() > 0 { + c := popRanked(s.heap) + // Build the follow-ups first, even for a path we're about to skip: a + // finished path still has paths that come after it, and skipping without + // expanding would lose them. + s.expand(c.payload) + if s.terminal[c.id] { + continue + } + return c.cand, true + } + return entity.CandidatePath{}, false +} + +// expand builds the paths that come just after this one. +// +// Treat the flip set as something that only moves forward. If the highest +// position flipped so far is i, there are two ways to move on: +// +// add i+1 to the set {0,1} -> {0,1,2} flip one more guess +// swap i for i+1 {0,1} -> {0,2} move the last flip along +// +// From the best path both moves come to the same thing, so it has the single +// child {0}. With three choices the full tree is: +// +// {} -> {0} +// {0} -> {0,1}, {1} +// {1} -> {1,2}, {2} +// {0,1} -> {0,1,2}, {0,2} +// {2}, {0,2}, {1,2}, {0,1,2} have nothing left to flip above position 2 +// +// That is all 8 sets, each appearing exactly once — and it has to be exactly +// once, because you can always tell which move produced a set: if its last two +// numbers are next to each other it was the first move, otherwise the second. +// Since there is only one way to reach any set, we never need to remember what +// we have already built. +// +// Both moves flip a guess no cheaper than the one they replace or build on, +// because choices is sorted cheapest-first. So a path never scores above the +// one it came from, which is what lets the heap hand paths out in score order. +func (s *headStream) expand(flipped flipSet) { + n := len(flipped) + last := -1 + if n > 0 { + last = flipped[n-1] + } + // Every choice past the last flip is already used up, so there is no move + // left to make and this path has no follow-ups. + if last+1 >= len(s.choices) { + return + } + + // Flip one more guess, keeping the ones already flipped. + extended := make(flipSet, n+1) + copy(extended, flipped) + extended[n] = last + 1 + s.push(extended) + + // Move the last flip one position along. The best path has no flip to move, + // so this second move doesn't apply to it. + if n > 0 { + advanced := make(flipSet, n) + copy(advanced, flipped) + advanced[n-1] = last + 1 + s.push(advanced) + } +} + +// push builds the path a flip set describes and puts it in the heap. It starts +// from the head's bets, switches each flipped choice to its alternative, and +// multiplies out the score along the way. Dependencies that already finished +// are certainties, so they don't change the score. +func (s *headStream) push(flipped flipSet) { + // Work on a copy. s.bets is the shared template every path starts from, so + // writing a flip straight into it would corrupt every path built after this + // one. + bets := make([]entity.DependencyBet, len(s.bets)) + copy(bets, s.bets) + + // Walk the choices and the flipped positions together. Both are in + // increasing order, so one cursor into flipped is enough to tell whether the + // choice we're looking at is one of the flipped ones. + score := 1.0 + next := 0 + for i, c := range s.choices { + if next < len(flipped) && flipped[next] == i { + // Flipped: take the other bet and the worse chance. + bets[c.slot] = entity.DependencyBet{Batch: c.dep, Bet: c.alternativeBet} + score *= c.alternative + next++ + continue + } + // Not flipped: the template already holds the preferred bet, so only the + // score needs the chance folded in. + score *= c.preferred + } + + path := entity.SpeculationPath{Head: s.head, Bets: bets} + s.generated++ + pushRanked(s.heap, combination{ + cand: entity.CandidatePath{Path: path, RankingScore: score}, + // Hash the path once here rather than every time the heap compares it. + id: path.ID(), + payload: flipped, + }) +} + +// iterator merges the per-head streams into one best-first stream. +// +// The heap holds each head's current offer — a real path, already built, never +// a guess at what a head might produce. That is what makes taking the top of +// the heap correct: a head can't be hiding anything better behind the path it +// is currently offering, so the best offer on the table is the best path in the +// queue. +type iterator struct { + // streams is every head stream we set up. The merge itself runs off the + // heap; this is kept because the iterator owns these streams, and tests read + // it to check how much work each head has done. + streams []*headStream + heap *rankedHeap[*headStream] +} + +// queueNext asks a stream for its next path and puts it on the heap as that +// head's offer. A stream with nothing left simply drops out of the merge. +func (it *iterator) queueNext(s *headStream) { + c, ok := s.Next() + if !ok { + return + } + pushRanked(it.heap, front{cand: c, id: c.Path.ID(), payload: s}) +} + +// Next hands out the best path across all heads, then asks that head for its +// next one to replace it. +func (it *iterator) Next(_ context.Context) (entity.CandidatePath, bool, error) { + if it.heap.Len() == 0 { + return entity.CandidatePath{}, false, nil + } + f := popRanked(it.heap) + it.queueNext(f.payload) + return f.cand, true, nil +} + +// terminalPathIDs collects, for each head, the IDs of its paths that already +// finished — passed, failed, or cancelled — so we don't offer them again. +func terminalPathIDs(pathSets []entity.SpeculationPathSet) map[string]map[string]bool { + byHead := make(map[string]map[string]bool, len(pathSets)) + for _, set := range pathSets { + for _, entry := range set.Paths { + if !entry.Status.IsTerminal() { + continue + } + // Only build a set for heads that actually have finished paths; most + // heads have none, and a missing entry reads as empty later. + ids := byHead[set.BatchID] + if ids == nil { + ids = make(map[string]bool) + byHead[set.BatchID] = ids + } + ids[entry.ID] = true + } + } + return byHead +} diff --git a/submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go b/submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go new file mode 100644 index 00000000..56d91568 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go @@ -0,0 +1,876 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// 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 bestfirst + +import ( + "context" + "fmt" + "math/rand" + "sort" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/scorer" + "github.com/uber/submitqueue/submitqueue/extension/speculation/generator" +) + +// stubScorer scores each batch by ID, defaulting to 0.5 for unknown batches. It +// is a minimal scorer.Scorer for exercising the generator without a resolver. +type stubScorer struct { + scores map[string]float64 +} + +func (s stubScorer) Score(_ context.Context, b entity.Batch) (float64, error) { + if v, ok := s.scores[b.ID]; ok { + return v, nil + } + return 0.5, nil +} + +func scored(scores map[string]float64) scorer.Scorer { + return stubScorer{scores: scores} +} + +// drainAll pulls every candidate from an iterator. +func drainAll(t *testing.T, iter generator.PathIterator) []entity.CandidatePath { + t.Helper() + var out []entity.CandidatePath + for { + c, ok, err := iter.Next(context.Background()) + require.NoError(t, err) + if !ok { + break + } + out = append(out, c) + } + return out +} + +// forHead returns only the candidates whose head is headID. +func forHead(cands []entity.CandidatePath, headID string) []entity.CandidatePath { + var out []entity.CandidatePath + for _, c := range cands { + if c.Path.Head == headID { + out = append(out, c) + } + } + return out +} + +// betFor returns the bet a path carries for a given dependency batch. +func betFor(p entity.SpeculationPath, dep string) entity.DependencyBetType { + for _, b := range p.Bets { + if b.Batch == dep { + return b.Bet + } + } + return entity.BetUnknown +} + +// betKey renders a path's bets in queue order, to compare whole combinations. +func betKey(p entity.SpeculationPath) string { + var b strings.Builder + for _, bet := range p.Bets { + b.WriteString(bet.Batch) + b.WriteByte('=') + b.WriteString(string(bet.Bet)) + b.WriteByte(';') + } + return b.String() +} + +// pathIDs renders each candidate's path ID, in order. +func pathIDs(cands []entity.CandidatePath) []string { + out := make([]string, 0, len(cands)) + for _, c := range cands { + out = append(out, c.Path.ID()) + } + return out +} + +// streamFor reaches into the iterator for one head's stream, so the laziness +// invariants can be asserted on the combinations it has actually materialized. +func streamFor(t *testing.T, iter generator.PathIterator, headID string) *headStream { + t.Helper() + it, ok := iter.(*iterator) + require.True(t, ok, "iterator is not the bestfirst merge") + for _, s := range it.streams { + if s.head == headID { + return s + } + } + t.Fatalf("no stream for head %q", headID) + return nil +} + +// countingScorer records how many times each batch is scored. +type countingScorer struct { + scores map[string]float64 + calls map[string]int + total int +} + +func newCountingScorer(scores map[string]float64) *countingScorer { + return &countingScorer{scores: scores, calls: map[string]int{}} +} + +func (c *countingScorer) Score(_ context.Context, b entity.Batch) (float64, error) { + c.calls[b.ID]++ + c.total++ + if v, ok := c.scores[b.ID]; ok { + return v, nil + } + return 0.5, nil +} + +// wideHead builds one Speculating head over n unresolved dependencies, each at a +// distinct landing chance so no two combinations tie. +func wideHead(n int) ([]entity.Batch, scorer.Scorer) { + head := entity.Batch{ID: "q/head", State: entity.BatchStateSpeculating} + batches := []entity.Batch{} + scores := map[string]float64{} + for i := 0; i < n; i++ { + dep := fmt.Sprintf("q/d%02d", i) + head.Dependencies = append(head.Dependencies, dep) + // Created deps are unresolved (so they're scored) but not eligible heads. + batches = append(batches, entity.Batch{ID: dep, State: entity.BatchStateCreated}) + scores[dep] = 0.55 + 0.02*float64(i) + } + return append(batches, head), scored(scores) +} + +func TestBestFirst_OrderingAndEnumeration(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSpeculating}, + {ID: "q/2", State: entity.BatchStateSpeculating}, + {ID: "q/3", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1", "q/2"}}, + } + sc := scored(map[string]float64{"q/1": 0.9, "q/2": 0.8}) + + iter, err := New(2, sc).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/3") + + // Two unresolved dependencies -> 2^2 candidates. + require.Len(t, cands, 4) + + // Best-first: the all-included path leads, scoring the product of landing + // chances (0.9 * 0.8), and scores descend from there. + assert.Equal(t, entity.BetIncluded, betFor(cands[0].Path, "q/1")) + assert.Equal(t, entity.BetIncluded, betFor(cands[0].Path, "q/2")) + assert.InDelta(t, 0.72, cands[0].RankingScore, 1e-9) + for i := 1; i < len(cands); i++ { + assert.LessOrEqual(t, cands[i].RankingScore, cands[i-1].RankingScore) + } + // The least optimistic candidate excludes both dependencies: (1-0.9)(1-0.8). + last := cands[len(cands)-1] + assert.Equal(t, entity.BetExcluded, betFor(last.Path, "q/1")) + assert.Equal(t, entity.BetExcluded, betFor(last.Path, "q/2")) + assert.InDelta(t, 0.02, last.RankingScore, 1e-9) +} + +func TestBestFirst_PinsResolvedDependencies(t *testing.T) { + t.Run("landed dependency forced included", func(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSucceeded}, + {ID: "q/2", State: entity.BatchStateSpeculating}, + {ID: "q/3", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1", "q/2"}}, + } + iter, err := New(2, scored(map[string]float64{"q/2": 0.8})).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/3") + + // q/1 is pinned, only q/2 varies -> 2 candidates. + require.Len(t, cands, 2) + for _, c := range cands { + assert.Equal(t, entity.BetIncluded, betFor(c.Path, "q/1")) + } + }) + + t.Run("failed dependency forced excluded", func(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateFailed}, + {ID: "q/2", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1"}}, + } + iter, err := New(2, scored(nil)).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/2") + + require.Len(t, cands, 1) + assert.Equal(t, entity.BetExcluded, betFor(cands[0].Path, "q/1")) + }) + + t.Run("pinned dependencies drop out of the search", func(t *testing.T) { + // Three dependencies, but two are resolved facts: only q/open is still a + // guess, so the head yields 2 paths rather than 8. + batches := []entity.Batch{ + {ID: "q/landed", State: entity.BatchStateSucceeded}, + {ID: "q/dropped", State: entity.BatchStateFailed}, + {ID: "q/open", State: entity.BatchStateCreated}, + {ID: "q/H", State: entity.BatchStateSpeculating, + Dependencies: []string{"q/landed", "q/dropped", "q/open"}}, + } + iter, err := New(3, scored(map[string]float64{"q/open": 0.7})). + Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/H") + + require.Len(t, cands, 2) + // Pinned bets are identical on both paths and stay in queue order; the + // score reflects only the open dependency, since a fact is a certainty. + assert.Equal(t, "q/landed=included;q/dropped=excluded;q/open=included;", betKey(cands[0].Path)) + assert.InDelta(t, 0.7, cands[0].RankingScore, 1e-9) + assert.Equal(t, "q/landed=included;q/dropped=excluded;q/open=excluded;", betKey(cands[1].Path)) + assert.InDelta(t, 0.3, cands[1].RankingScore, 1e-9) + }) + + t.Run("cancelled dependency forced excluded", func(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateCancelled}, + {ID: "q/2", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1"}}, + } + iter, err := New(2, scored(nil)).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/2") + + require.Len(t, cands, 1) + assert.Equal(t, entity.BetExcluded, betFor(cands[0].Path, "q/1")) + }) +} + +func TestBestFirst_SkipsTerminalPaths(t *testing.T) { + // A path is passed over only once its build reached a terminal state. A path + // still in flight stays a candidate — whether to re-propose it is the + // Speculator's call, not the generator's. + tests := []struct { + name string + status entity.SpeculationPathStatus + skipped bool + }{ + {name: "passed", status: entity.SpeculationPathStatusPassed, skipped: true}, + {name: "failed", status: entity.SpeculationPathStatusFailed, skipped: true}, + {name: "cancelled", status: entity.SpeculationPathStatusCancelled, skipped: true}, + {name: "pending is not terminal", status: entity.SpeculationPathStatusPending}, + {name: "building is not terminal", status: entity.SpeculationPathStatusBuilding}, + {name: "cancelling is not terminal", status: entity.SpeculationPathStatusCancelling}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSpeculating}, + {ID: "q/2", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1"}}, + } + + // The leading (included) path for q/2 carries the status under test. + included := entity.SpeculationPath{ + Head: "q/2", + Bets: []entity.DependencyBet{{Batch: "q/1", Bet: entity.BetIncluded}}, + } + pathSets := []entity.SpeculationPathSet{{ + BatchID: "q/2", + Paths: []entity.SpeculationPathEntry{{ + ID: included.ID(), + Path: included, + Status: tt.status, + }}, + }} + + iter, err := New(2, scored(map[string]float64{"q/1": 0.7})). + Open(context.Background(), batches, pathSets) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/2") + + if tt.skipped { + // Only the excluded candidate remains. + require.Len(t, cands, 1) + assert.Equal(t, entity.BetExcluded, betFor(cands[0].Path, "q/1")) + return + } + require.Len(t, cands, 2) + assert.Equal(t, entity.BetIncluded, betFor(cands[0].Path, "q/1")) + }) + } +} + +func TestBestFirst_EmitsExactSequenceAcrossHeads(t *testing.T) { + // q/solo has nothing to wait on; q/pair bets on two in-flight batches. The + // two heads interleave by score rather than draining one at a time. + batches := []entity.Batch{ + {ID: "q/d0", State: entity.BatchStateCreated}, + {ID: "q/d1", State: entity.BatchStateCreated}, + {ID: "q/solo", State: entity.BatchStateSpeculating}, + {ID: "q/pair", State: entity.BatchStateSpeculating, Dependencies: []string{"q/d0", "q/d1"}}, + } + sc := scored(map[string]float64{"q/d0": 0.9, "q/d1": 0.8}) + + iter, err := New(2, sc).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := drainAll(t, iter) + + want := []struct { + head string + bets string + score float64 + }{ + {head: "q/solo", bets: "", score: 1.0}, + {head: "q/pair", bets: "q/d0=included;q/d1=included;", score: 0.72}, + {head: "q/pair", bets: "q/d0=included;q/d1=excluded;", score: 0.18}, + {head: "q/pair", bets: "q/d0=excluded;q/d1=included;", score: 0.08}, + {head: "q/pair", bets: "q/d0=excluded;q/d1=excluded;", score: 0.02}, + } + require.Len(t, cands, len(want)) + for i, w := range want { + assert.Equal(t, w.head, cands[i].Path.Head, "head at %d", i) + assert.Equal(t, w.bets, betKey(cands[i].Path), "bets at %d", i) + assert.InDelta(t, w.score, cands[i].RankingScore, 1e-9, "score at %d", i) + } +} + +func TestBestFirst_PreferredBetFollowsLikelierOutcome(t *testing.T) { + // q/unlikely is likelier to fail than to land, so the head's best path bets + // against it. The leading path is the preferred-bet path, not the + // all-included one. + batches := []entity.Batch{ + {ID: "q/likely", State: entity.BatchStateCreated}, + {ID: "q/unlikely", State: entity.BatchStateCreated}, + {ID: "q/H", State: entity.BatchStateSpeculating, + Dependencies: []string{"q/likely", "q/unlikely"}}, + } + sc := scored(map[string]float64{"q/likely": 0.8, "q/unlikely": 0.3}) + + iter, err := New(2, sc).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := drainAll(t, iter) + + want := []struct { + bets string + score float64 + }{ + {bets: "q/likely=included;q/unlikely=excluded;", score: 0.56}, + {bets: "q/likely=included;q/unlikely=included;", score: 0.24}, + {bets: "q/likely=excluded;q/unlikely=excluded;", score: 0.14}, + {bets: "q/likely=excluded;q/unlikely=included;", score: 0.06}, + } + require.Len(t, cands, len(want)) + for i, w := range want { + assert.Equal(t, w.bets, betKey(cands[i].Path), "bets at %d", i) + assert.InDelta(t, w.score, cands[i].RankingScore, 1e-9, "score at %d", i) + } + + // Spelled out: the all-included path exists but only ranks second. + assert.Equal(t, entity.BetExcluded, betFor(cands[0].Path, "q/unlikely")) +} + +func TestBestFirst_SkipsHeadOverDepthBound(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSpeculating}, + {ID: "q/2", State: entity.BatchStateSpeculating}, + {ID: "q/3", State: entity.BatchStateSpeculating}, + {ID: "q/4", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1", "q/2", "q/3"}}, + } + + iter, err := New(2, scored(nil)).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/4") + + // Three unresolved dependencies exceed the depth bound of 2. + assert.Empty(t, cands) +} + +func TestBestFirst_SkipsHaltedHeads(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSucceeded}, + {ID: "q/2", State: entity.BatchStateCancelling, Dependencies: []string{"q/1"}}, + {ID: "q/3", State: entity.BatchStateFailed, Dependencies: []string{"q/1"}}, + } + + iter, err := New(2, scored(nil)).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := drainAll(t, iter) + + // Succeeded, cancelling, and failed heads all produce no candidates. + assert.Empty(t, cands) +} + +func TestBestFirst_HeadWithNoDependencies(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSpeculating}, + } + + iter, err := New(2, scored(nil)).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := drainAll(t, iter) + + require.Len(t, cands, 1) + assert.Equal(t, "q/1", cands[0].Path.Head) + assert.Empty(t, cands[0].Path.Bets) + assert.InDelta(t, 1.0, cands[0].RankingScore, 1e-9) +} + +func TestBestFirst_InterleavesHeadsBestFirst(t *testing.T) { + batches := []entity.Batch{ + // q/1: no deps -> single candidate, score 1.0. + {ID: "q/1", State: entity.BatchStateSpeculating}, + // q/2: one unresolved dep scored 0.6 -> included (0.6) and excluded (0.4). + {ID: "q/dep", State: entity.BatchStateSpeculating}, + {ID: "q/2", State: entity.BatchStateSpeculating, Dependencies: []string{"q/dep"}}, + } + + iter, err := New(2, scored(map[string]float64{"q/dep": 0.6})).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := drainAll(t, iter) + + // Scores are globally descending across all heads. + for i := 1; i < len(cands); i++ { + assert.LessOrEqual(t, cands[i].RankingScore, cands[i-1].RankingScore) + } + // The two dep-free heads (q/1, q/dep) lead at score 1.0. + assert.InDelta(t, 1.0, cands[0].RankingScore, 1e-9) +} + +func TestBestFirst_PropagatesScorerError(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSpeculating}, + {ID: "q/2", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1"}}, + } + + iter, err := New(2, errScorer{}).Open(context.Background(), batches, nil) + assert.Error(t, err) + assert.Nil(t, iter) +} + +// errScorer always fails, to exercise error propagation from scoring. +type errScorer struct{} + +func (errScorer) Score(context.Context, entity.Batch) (float64, error) { + return 0, assert.AnError +} + +// constScorer scores every batch identically, regardless of ID. +type constScorer struct{ v float64 } + +func (c constScorer) Score(context.Context, entity.Batch) (float64, error) { return c.v, nil } + +func TestBestFirst_SkipsNonSpeculatingHeads(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/created", State: entity.BatchStateCreated}, + {ID: "q/merging", State: entity.BatchStateMerging}, + {ID: "q/speculating", State: entity.BatchStateSpeculating}, + } + + iter, err := New(2, scored(nil)).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := drainAll(t, iter) + + // Created and Merging heads supply facts but are never action targets; only + // the Speculating head produces a candidate. + require.Len(t, cands, 1) + assert.Equal(t, "q/speculating", cands[0].Path.Head) +} + +func TestBestFirst_CrossHeadOrderingWithTerminalFilteredTop(t *testing.T) { + // Head A's best (all-included) path is terminal-filtered, so its real front + // (0.1) ranks below head B (0.5) even though A's seed (0.9) ranks above B. + // The merge must emit B before A's real candidate. + batches := []entity.Batch{ + {ID: "q/A", State: entity.BatchStateSpeculating, Dependencies: []string{"q/depA"}}, + {ID: "q/B", State: entity.BatchStateSpeculating, Dependencies: []string{"q/depB"}}, + // Deps are Created: unresolved (so they're scored) but not eligible heads. + {ID: "q/depA", State: entity.BatchStateCreated}, + {ID: "q/depB", State: entity.BatchStateCreated}, + } + aIncluded := entity.SpeculationPath{ + Head: "q/A", + Bets: []entity.DependencyBet{{Batch: "q/depA", Bet: entity.BetIncluded}}, + } + pathSets := []entity.SpeculationPathSet{{ + BatchID: "q/A", + Paths: []entity.SpeculationPathEntry{{ID: aIncluded.ID(), Path: aIncluded, Status: entity.SpeculationPathStatusPassed}}, + }} + + sc := scored(map[string]float64{"q/depA": 0.9, "q/depB": 0.5}) + iter, err := New(2, sc).Open(context.Background(), batches, pathSets) + require.NoError(t, err) + cands := drainAll(t, iter) + + // B (both combos at 0.5) precedes A's real front; A contributes only its + // excluded path (0.1) since the included one was terminal-filtered. + require.Len(t, cands, 3) + assert.Equal(t, "q/B", cands[0].Path.Head) + last := cands[len(cands)-1] + assert.Equal(t, "q/A", last.Path.Head) + assert.Equal(t, entity.BetExcluded, betFor(last.Path, "q/depA")) + assert.InDelta(t, 0.1, last.RankingScore, 1e-9) +} + +func TestBestFirst_AbsentDependencyFallsBackToCoinFlip(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/H", State: entity.BatchStateSpeculating, Dependencies: []string{"q/missing"}}, + } + + // The scorer would return 0.9, but q/missing is absent from the batches, so + // the generator falls back to a 0.5 coin flip without consulting the scorer. + iter, err := New(2, constScorer{0.9}).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/H") + + require.Len(t, cands, 2) + for _, c := range cands { + assert.InDelta(t, 0.5, c.RankingScore, 1e-9) + } +} + +func TestBestFirst_DepthBoundZero(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/nodep", State: entity.BatchStateSpeculating}, + {ID: "q/dep", State: entity.BatchStateCreated}, + {ID: "q/withdep", State: entity.BatchStateSpeculating, Dependencies: []string{"q/dep"}}, + } + + iter, err := New(0, scored(nil)).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := drainAll(t, iter) + + // depthBound 0 admits only the dependency-free Speculating head. + require.Len(t, cands, 1) + assert.Equal(t, "q/nodep", cands[0].Path.Head) +} + +func TestBestFirst_GeneratesOnlyWhatIsPulled(t *testing.T) { + // 12 unresolved dependencies is an outcome space of 4096 paths. + const deps, space = 12, 1 << 12 + batches, sc := wideHead(deps) + + iter, err := New(deps, sc).Open(context.Background(), batches, nil) + require.NoError(t, err) + s := streamFor(t, iter, "q/head") + + // Open seats the head's best path and nothing more: one combination popped, + // at most two children pushed. + require.LessOrEqual(t, s.generated, 3) + + for i := 0; i < 3; i++ { + _, ok, err := iter.Next(context.Background()) + require.NoError(t, err) + require.True(t, ok) + } + + // Each pull expands by at most two combinations, so four pops in total can + // have materialized at most 1+2*4 paths — a tiny fraction of the space. + assert.LessOrEqual(t, s.generated, 9) + assert.Less(t, s.generated, space) +} + +func TestBestFirst_HeadAdvancesOnlyWhenConsumed(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/a0", State: entity.BatchStateCreated}, + {ID: "q/a1", State: entity.BatchStateCreated}, + {ID: "q/a2", State: entity.BatchStateCreated}, + {ID: "q/b0", State: entity.BatchStateCreated}, + {ID: "q/b1", State: entity.BatchStateCreated}, + {ID: "q/b2", State: entity.BatchStateCreated}, + {ID: "q/A", State: entity.BatchStateSpeculating, Dependencies: []string{"q/a0", "q/a1", "q/a2"}}, + {ID: "q/B", State: entity.BatchStateSpeculating, Dependencies: []string{"q/b0", "q/b1", "q/b2"}}, + } + // A's best path (0.9^3) outranks B's (0.6^3), so A is consumed first. + sc := scored(map[string]float64{ + "q/a0": 0.9, "q/a1": 0.9, "q/a2": 0.9, + "q/b0": 0.6, "q/b1": 0.6, "q/b2": 0.6, + }) + + iter, err := New(3, sc).Open(context.Background(), batches, nil) + require.NoError(t, err) + a, b := streamFor(t, iter, "q/A"), streamFor(t, iter, "q/B") + + // Seating a head materializes its best path and that path's single child — + // not the eight paths in its space. + require.Equal(t, 2, a.generated) + require.Equal(t, 2, b.generated) + + // Consuming A's candidate advances A only; B's stream is untouched. + c, ok, err := iter.Next(context.Background()) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, "q/A", c.Path.Head) + assert.Greater(t, a.generated, 2) + assert.Equal(t, 2, b.generated) +} + +func TestBestFirst_DrainYieldsEveryCombinationOnce(t *testing.T) { + deps := []string{"q/d0", "q/d1", "q/d2"} + batches := []entity.Batch{ + {ID: "q/d0", State: entity.BatchStateCreated}, + {ID: "q/d1", State: entity.BatchStateCreated}, + {ID: "q/d2", State: entity.BatchStateCreated}, + {ID: "q/head", State: entity.BatchStateSpeculating, Dependencies: deps}, + } + sc := scored(map[string]float64{"q/d0": 0.9, "q/d1": 0.7, "q/d2": 0.6}) + + iter, err := New(3, sc).Open(context.Background(), batches, nil) + require.NoError(t, err) + cands := drainAll(t, iter) + + require.Len(t, cands, 8) + + // Every include/exclude combination appears exactly once, and each path + // carries its bets in queue order. + seen := map[string]int{} + for _, c := range cands { + got := make([]string, 0, len(c.Path.Bets)) + for _, bet := range c.Path.Bets { + got = append(got, bet.Batch) + } + assert.Equal(t, deps, got, "bets must stay in dependency order") + seen[betKey(c.Path)]++ + } + for mask := 0; mask < 8; mask++ { + var want strings.Builder + for i, dep := range deps { + bet := entity.BetExcluded + if mask&(1<