Visitar URL original
feat(bigtable): add AFE picker (Simple / LeastInFlight / LeastLatency… · googleapis/google-cloud-go@bcbf714 · GitHub
Skip to content

Commit bcbf714

Browse files
sushanbsushanb
authored
feat(bigtable): add AFE picker (Simple / LeastInFlight / LeastLatency) (#20204)
## Summary Adds the AFE (Application Front End) picker interface and three implementations, plus the two consumer-facing types they operate on. Pure function over `[]afeSnapshot` — no ties to a `sessionList` producer, no lock ownership, no pool integration. The producer + wiring land in follow-up PRs. - **`afe_picker.go`** — `AfePicker` interface + `SimpleAfePicker` (uniform random), `LeastInFlightAfePicker` (K-choice min NumOutstanding), `LeastLatencyAfePicker` (K-choice min e2e PeakEwma). Every `PickAfe` returns a `PickDecision{Candidates, Winner, Reason}` so debug surfaces can trace picker reasoning without re-running the pick. - **`afe_picker_test.go`** — 10 unit tests: empty-input handling per picker, uniform-random distribution smoke, min-cost correctness, K-choice sub-sampling, in-place mutation of caller's slice, and picker-name identity. - **`afe_types.go`** — `afeID int64` + `afeSnapshot` struct (5 fields, 3 used by the picker today). Split into its own file so the picker ships as a pure function over snapshot slices; the producer (`sessionList`) lands in a follow-up PR. `kChoiceMinCost` implements partial-Fisher-Yates over the caller's slice in place — a defensive copy per pick cost ~4µs at steady-state QPS, so callers own the slice's lifetime. ## Test plan - [x] `go test ./bigtable/internal/transport/ -run "AfePicker|KChoice|Decision" -count=1 -short` → 10/10 pass locally. - [x] `go build ./bigtable/internal/transport/` clean. - [ ] CI green.
1 parent 6e776eb commit bcbf714

3 files changed

Lines changed: 424 additions & 0 deletions

File tree

‎bigtable/internal/transport/afe_picker.go‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package internal
16+
17+
import (
18+
"math/rand/v2"
19+
)
20+
21+
// PickCandidate is one AFE the picker considered during a K-choice draw,
22+
// with the cost value the picker's decision rule used to score it.
23+
// Cost's interpretation depends on the picker in play: NumOutstanding
24+
// (int as float) for LeastInFlight, e2e PeakEwma nanos for LeastLatency,
25+
// 0 for Simple (which ignores cost).
26+
type PickCandidate struct {
27+
AfeID AfeID
28+
Cost float64
29+
}
30+
31+
// PickDecision captures what candidates a picker sampled, which one won,
32+
// and why. Populated on every PickAfe call so operators can trace picker
33+
// reasoning through the debug surface without re-running the pick.
34+
type PickDecision struct {
35+
// Candidates is the K sampled AFEs (K == 1 for SimplePicker,
36+
// otherwise K == min(RandomSubsetSize, len(ready)) via partial
37+
// Fisher-Yates). Empty when no ready AFE existed.
38+
Candidates []PickCandidate
39+
// Winner is the AFE the picker returned. Zero when ready was empty.
40+
Winner AfeID
41+
// Reason is a short lower-kebab tag identifying the decision rule
42+
// ("uniform-random" / "min-inflight" / "min-latency" /
43+
// "no-candidates"). Machine-readable; consumers map to prose.
44+
Reason string
45+
}
46+
47+
// AfePicker picks one AFE from a snapshot of ready buckets AND returns
48+
// the decision metadata for the debug surface. The winner travels back
49+
// into the producer's Checkout by id — no producer-owned pointer ever
50+
// leaves the producer's lock. The PickDecision is passed to the pool's
51+
// pick-decision recorder.
52+
//
53+
// picked == false means "no AFE eligible" — the pool treats that the
54+
// same as len(ready) == 0 (park the caller, kick scale-up).
55+
//
56+
// Implementations MAY mutate ready in place; callers must pass a
57+
// throwaway slice.
58+
type AfePicker interface {
59+
PickAfe(ready []AfeSnapshot) (winner AfeID, picked bool, decision PickDecision)
60+
Name() string
61+
}
62+
63+
// SimpleAfePicker chooses an AFE uniformly at random from the ready set.
64+
type SimpleAfePicker struct{}
65+
66+
// NewSimpleAfePicker constructs a SimpleAfePicker.
67+
func NewSimpleAfePicker() *SimpleAfePicker { return &SimpleAfePicker{} }
68+
69+
// Name returns "simple".
70+
func (SimpleAfePicker) Name() string { return "simple" }
71+
72+
// PickAfe uniformly-at-random picks one bucket from ready.
73+
func (SimpleAfePicker) PickAfe(ready []AfeSnapshot) (AfeID, bool, PickDecision) {
74+
if len(ready) == 0 {
75+
return 0, false, PickDecision{Reason: "no-candidates"}
76+
}
77+
winner := ready[rand.IntN(len(ready))]
78+
return winner.ID, true, PickDecision{
79+
Candidates: []PickCandidate{{AfeID: winner.ID, Cost: 0}},
80+
Winner: winner.ID,
81+
Reason: "uniform-random",
82+
}
83+
}
84+
85+
// LeastInFlightAfePicker picks the AFE with the smallest in-flight count.
86+
// Draws K distinct candidates via partial Fisher-Yates over the ready
87+
// snapshot and returns the min-cost one. RandomSubsetSize caps K; when
88+
// it's <=0 or >= len(ready) every candidate is considered.
89+
type LeastInFlightAfePicker struct {
90+
// RandomSubsetSize caps the K-choice draw. 0 or negative means
91+
// "consider all candidates".
92+
RandomSubsetSize int
93+
}
94+
95+
// NewLeastInFlightAfePicker constructs a LeastInFlightAfePicker.
96+
func NewLeastInFlightAfePicker(randomSubsetSize int) *LeastInFlightAfePicker {
97+
return &LeastInFlightAfePicker{RandomSubsetSize: randomSubsetSize}
98+
}
99+
100+
// Name returns "least-inflight".
101+
func (LeastInFlightAfePicker) Name() string { return "least-inflight" }
102+
103+
// PickAfe returns the AFE with the fewest NumOutstanding among K
104+
// randomly-drawn ready candidates.
105+
func (p LeastInFlightAfePicker) PickAfe(ready []AfeSnapshot) (AfeID, bool, PickDecision) {
106+
winner, picked, cands := kChoiceMinCost(ready, p.RandomSubsetSize, func(s AfeSnapshot) float64 {
107+
return float64(s.NumOutstanding)
108+
})
109+
return decisionFor(winner, picked, cands, "min-inflight")
110+
}
111+
112+
// LeastLatencyAfePicker picks the AFE with the lowest per-AFE e2e
113+
// PeakEwma cost. Same K-choice partial Fisher-Yates as
114+
// LeastInFlightAfePicker.
115+
type LeastLatencyAfePicker struct {
116+
RandomSubsetSize int
117+
}
118+
119+
// NewLeastLatencyAfePicker constructs a LeastLatencyAfePicker.
120+
func NewLeastLatencyAfePicker(randomSubsetSize int) *LeastLatencyAfePicker {
121+
return &LeastLatencyAfePicker{RandomSubsetSize: randomSubsetSize}
122+
}
123+
124+
// Name returns "least-latency".
125+
func (LeastLatencyAfePicker) Name() string { return "least-latency" }
126+
127+
// PickAfe returns the AFE with the smallest E2eCost among K randomly-
128+
// drawn ready candidates.
129+
func (p LeastLatencyAfePicker) PickAfe(ready []AfeSnapshot) (AfeID, bool, PickDecision) {
130+
winner, picked, cands := kChoiceMinCost(ready, p.RandomSubsetSize, func(s AfeSnapshot) float64 {
131+
return s.E2eCost
132+
})
133+
return decisionFor(winner, picked, cands, "min-latency")
134+
}
135+
136+
// decisionFor packages kChoiceMinCost's return into a PickDecision.
137+
func decisionFor(winner AfeID, picked bool, cands []PickCandidate, reason string) (AfeID, bool, PickDecision) {
138+
if !picked {
139+
return 0, false, PickDecision{Reason: "no-candidates"}
140+
}
141+
return winner, true, PickDecision{
142+
Candidates: cands,
143+
Winner: winner,
144+
Reason: reason,
145+
}
146+
}
147+
148+
// kChoiceMinCost implements partial-Fisher-Yates + min-cost selection
149+
// over a snapshot slice. K is clamped to len(ready); K<=0 means scan
150+
// every candidate.
151+
//
152+
// The algorithm mutates ready in place (swap-to-front). Callers must
153+
// pass a throwaway slice — every production call site produces one via
154+
// the producer's snapshot method, which allocates a fresh copy per call.
155+
// cost is invoked at most K times. Returns the winner plus the list of
156+
// sampled candidates (with their costs) so callers can build a
157+
// PickDecision for the debug surface.
158+
//
159+
// The previous implementation allocated a defensive copy of ready per
160+
// call; profiling showed it costing ~4µs at the workload's steady-state
161+
// QPS since the picker runs on every CheckoutSession. Removed because
162+
// the caller doesn't need ready preserved.
163+
func kChoiceMinCost(ready []AfeSnapshot, k int, cost func(AfeSnapshot) float64) (AfeID, bool, []PickCandidate) {
164+
n := len(ready)
165+
if n == 0 {
166+
return 0, false, nil
167+
}
168+
if k <= 0 || k > n {
169+
k = n
170+
}
171+
172+
sampled := make([]PickCandidate, 0, k)
173+
var best AfeID
174+
haveBest := false
175+
bestCost := -1.0
176+
for i := 0; i < k; i++ {
177+
j := i + rand.IntN(n-i)
178+
s := ready[j]
179+
c := cost(s)
180+
sampled = append(sampled, PickCandidate{AfeID: s.ID, Cost: c})
181+
if !haveBest || c < bestCost {
182+
bestCost = c
183+
best = s.ID
184+
haveBest = true
185+
}
186+
ready[i], ready[j] = ready[j], ready[i]
187+
}
188+
return best, haveBest, sampled
189+
}

‎bigtable/internal/transport/afe_picker_test.go‎

Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package internal
16+
17+
import (
18+
"testing"
19+
)
20+
21+
// makeSnapshot builds an AfeSnapshot pointing at a fresh afeHandle with
22+
// the given id. Used only in these tests; production callers get
23+
// snapshots from sessionList.ReadyAfes.
24+
func makeSnapshot(id AfeID, inflight int, e2eCost float64) AfeSnapshot {
25+
return AfeSnapshot{
26+
ID: id,
27+
IdleCount: 1,
28+
NumOutstanding: inflight,
29+
E2eCost: e2eCost,
30+
}
31+
}
32+
33+
func TestSimpleAfePicker_EmptyReturnsNil(t *testing.T) {
34+
_, picked, dec := NewSimpleAfePicker().PickAfe(nil)
35+
if picked {
36+
t.Error("PickAfe(nil) picked = true, want false")
37+
}
38+
if dec.Reason != "no-candidates" {
39+
t.Errorf("PickAfe(nil) reason = %q, want no-candidates", dec.Reason)
40+
}
41+
}
42+
43+
func TestSimpleAfePicker_DistributesRoughly(t *testing.T) {
44+
snaps := []AfeSnapshot{
45+
makeSnapshot(1, 0, 0),
46+
makeSnapshot(2, 0, 0),
47+
makeSnapshot(3, 0, 0),
48+
}
49+
counts := map[AfeID]int{}
50+
p := NewSimpleAfePicker()
51+
const N = 3000
52+
for i := 0; i < N; i++ {
53+
id, _, _ := p.PickAfe(snaps)
54+
counts[id]++
55+
}
56+
for _, s := range snaps {
57+
got := counts[s.ID]
58+
// 3-way uniform expected value = 1000; allow ±25% (~750..1250).
59+
if got < 750 || got > 1250 {
60+
t.Errorf("AFE %d picked %d/%d times, want ~1000 ±25%%", s.ID, got, N)
61+
}
62+
}
63+
}
64+
65+
func TestLeastInFlightAfePicker_EmptyReturnsNil(t *testing.T) {
66+
_, picked, _ := NewLeastInFlightAfePicker(0).PickAfe(nil)
67+
if picked {
68+
t.Error("PickAfe(nil) picked = true, want false")
69+
}
70+
}
71+
72+
// When randomSubsetSize covers the full ready set, the min NumOutstanding
73+
// always wins deterministically.
74+
func TestLeastInFlightAfePicker_PicksMinWhenSubsetCoversAll(t *testing.T) {
75+
snaps := []AfeSnapshot{
76+
makeSnapshot(1, 5, 0),
77+
makeSnapshot(2, 1, 0), // min
78+
makeSnapshot(3, 3, 0),
79+
}
80+
p := NewLeastInFlightAfePicker(len(snaps))
81+
for i := 0; i < 100; i++ {
82+
id, _, _ := p.PickAfe(snaps)
83+
if id != 2 {
84+
t.Fatalf("iter %d: picked AFE %d, want 2", i, id)
85+
}
86+
}
87+
}
88+
89+
// K-choice-2 over 4 AFEs where one AFE has drastically lower inflight
90+
// should still pick that AFE the majority of the time.
91+
func TestLeastInFlightAfePicker_KChoicePrefersLowInflight(t *testing.T) {
92+
snaps := []AfeSnapshot{
93+
makeSnapshot(1, 100, 0),
94+
makeSnapshot(2, 100, 0),
95+
makeSnapshot(3, 1, 0), // the good one
96+
makeSnapshot(4, 100, 0),
97+
}
98+
p := NewLeastInFlightAfePicker(2)
99+
counts := map[AfeID]int{}
100+
const N = 4000
101+
for i := 0; i < N; i++ {
102+
id, _, _ := p.PickAfe(snaps)
103+
counts[id]++
104+
}
105+
// K=2 draws hit AFE 3 with prob 1 - (3/4)*(2/3) = 1/2; when drawn
106+
// it always wins. So expected share ≈ 50% ± sampling noise (±5%).
107+
share := float64(counts[3]) / float64(N)
108+
if share < 0.4 || share > 0.6 {
109+
t.Errorf("K=2 share for min-inflight AFE = %.3f, want ~0.5", share)
110+
}
111+
}
112+
113+
func TestLeastLatencyAfePicker_PicksMinCost(t *testing.T) {
114+
snaps := []AfeSnapshot{
115+
makeSnapshot(1, 0, 50.0),
116+
makeSnapshot(2, 0, 5.0), // min cost
117+
makeSnapshot(3, 0, 30.0),
118+
}
119+
p := NewLeastLatencyAfePicker(len(snaps))
120+
for i := 0; i < 100; i++ {
121+
id, _, _ := p.PickAfe(snaps)
122+
if id != 2 {
123+
t.Fatalf("iter %d: picked AFE %d, want 2 (min E2eCost)", i, id)
124+
}
125+
}
126+
}
127+
128+
func TestLeastLatencyAfePicker_EmptyReturnsNil(t *testing.T) {
129+
_, picked, _ := NewLeastLatencyAfePicker(0).PickAfe(nil)
130+
if picked {
131+
t.Error("PickAfe(nil) picked = true, want false")
132+
}
133+
}
134+
135+
// kChoiceMinCost mutates the caller's snapshot slice in place via
136+
// partial Fisher-Yates swap-to-front. Documented contract: callers must
137+
// own the slice. Production call site is CheckoutSession, which passes
138+
// p.sl.ReadyAfes() — a freshly-allocated slice per call. Verify the
139+
// mutation happens (all K positions get filled with distinct entries)
140+
// so a future re-introduction of a defensive copy is caught.
141+
func TestKChoiceMinCost_MutatesInputInPlace(t *testing.T) {
142+
snaps := []AfeSnapshot{
143+
makeSnapshot(1, 0, 0),
144+
makeSnapshot(2, 0, 0),
145+
makeSnapshot(3, 0, 0),
146+
}
147+
// After PickAfe with K=3, all three positions must contain the same
148+
// set of IDs (swap-to-front is a permutation, not a rewrite). Also
149+
// verify that the picker actually visited every slot: with K=len(snaps)
150+
// the first-K elements are exactly the K sampled candidates.
151+
NewLeastInFlightAfePicker(3).PickAfe(snaps)
152+
153+
seen := map[AfeID]bool{}
154+
for _, s := range snaps {
155+
seen[s.ID] = true
156+
}
157+
for _, want := range []AfeID{1, 2, 3} {
158+
if !seen[want] {
159+
t.Errorf("after PickAfe, ID %d missing from mutated slice (%v)", want, seen)
160+
}
161+
}
162+
}
163+
164+
// TestPickDecision_RecordsCandidatesAndReason verifies the K-choice
165+
// pickers surface every sampled candidate + the winning reason so
166+
// loadz can render the decision trace verbatim.
167+
func TestPickDecision_RecordsCandidatesAndReason(t *testing.T) {
168+
snaps := []AfeSnapshot{
169+
makeSnapshot(1, 5, 0),
170+
makeSnapshot(2, 1, 0),
171+
makeSnapshot(3, 3, 0),
172+
}
173+
_, _, dec := NewLeastInFlightAfePicker(2).PickAfe(snaps)
174+
if len(dec.Candidates) != 2 {
175+
t.Fatalf("Candidates len = %d, want 2 (K-choice-2)", len(dec.Candidates))
176+
}
177+
if dec.Reason != "min-inflight" {
178+
t.Errorf("Reason = %q, want min-inflight", dec.Reason)
179+
}
180+
if dec.Winner == 0 {
181+
t.Error("Winner is zero; expected the chosen AFE id")
182+
}
183+
184+
// LeastLatency uses e2e cost.
185+
_, _, dec2 := NewLeastLatencyAfePicker(len(snaps)).PickAfe(snaps)
186+
if dec2.Reason != "min-latency" {
187+
t.Errorf("LeastLatency reason = %q, want min-latency", dec2.Reason)
188+
}
189+
if len(dec2.Candidates) != len(snaps) {
190+
t.Errorf("LeastLatency Candidates len = %d, want %d (K covered all)", len(dec2.Candidates), len(snaps))
191+
}
192+
}
193+
194+
func TestAfePicker_Names(t *testing.T) {
195+
for _, tc := range []struct {
196+
p AfePicker
197+
want string
198+
}{
199+
{NewSimpleAfePicker(), "simple"},
200+
{NewLeastInFlightAfePicker(0), "least-inflight"},
201+
{NewLeastLatencyAfePicker(0), "least-latency"},
202+
} {
203+
if got := tc.p.Name(); got != tc.want {
204+
t.Errorf("%T.Name() = %q, want %q", tc.p, got, tc.want)
205+
}
206+
}
207+
}

0 commit comments

Comments
 (0)