Skip to content

Commit 84c3cd3

Browse files
sean-sean-
committed
refactor: pool score buffers and nil-guard connScoreSelect results
Replace per-call [8]float64 stack buffers with a sync.Pool that ratchets to the working set size. Eliminates heap escapes for the common case while handling arbitrarily large candidate sets. Add nil checks on connScoreSelect return values — if all candidates are skipped (warmup, overload), return an empty NextHop instead of dereferencing nil. Signed-off-by: Sean Chittenden <sean.chittenden@crowdstrike.com>
1 parent 48f5a65 commit 84c3cd3

4 files changed

Lines changed: 223 additions & 52 deletions

File tree

‎opensearchtransport/policy_doc_router.go‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -141,12 +141,9 @@ func (p *DocRouter) Eval(_ context.Context, req *http.Request) (NextHop, error)
141141

142142
shardCandidates, shardNum, shard := shardExactCandidates(p.cache.features, slot, effectiveRoutingKey, conns)
143143
if len(shardCandidates) > 0 {
144-
var scoresBuf [8]float64
145-
scores := scoresBuf[:len(shardCandidates)]
146-
if len(shardCandidates) > len(scoresBuf) {
147-
scores = make([]float64, len(shardCandidates))
148-
}
149-
best := connScoreSelect(shardCandidates, slot, shard, &shardCostForReads, "", loadPoolInfoReady(p.config.poolInfoReady), scores)
144+
scoreBuf := acquireScoreSlice(len(shardCandidates))
145+
best := connScoreSelect(shardCandidates, slot, shard, &shardCostForReads, "", loadPoolInfoReady(p.config.poolInfoReady), *scoreBuf)
146+
releaseScoreSlice(scoreBuf)
150147

151148
if obs := observerFromAtomic(&p.observer); obs != nil {
152149
key := indexName + "/" + docID
@@ -168,6 +165,9 @@ func (p *DocRouter) Eval(_ context.Context, req *http.Request) (NextHop, error)
168165
}))
169166
}
170167

168+
if best == nil {
169+
return NextHop{}, nil
170+
}
171171
return NextHop{Conn: best}, nil
172172
}
173173

@@ -192,12 +192,9 @@ func (p *DocRouter) Eval(_ context.Context, req *http.Request) (NextHop, error)
192192
slot.updateSmoothedMaxBucket(float64(maxBucket))
193193

194194
// Select best candidate with warmup-aware skip/accept.
195-
var scoresBuf [8]float64
196-
scores := scoresBuf[:len(candidates)]
197-
if len(candidates) > len(scoresBuf) {
198-
scores = make([]float64, len(candidates))
199-
}
200-
best := connScoreSelect(candidates, slot, nil, &shardCostForReads, "", loadPoolInfoReady(p.config.poolInfoReady), scores)
195+
scoreBuf := acquireScoreSlice(len(candidates))
196+
best := connScoreSelect(candidates, slot, nil, &shardCostForReads, "", loadPoolInfoReady(p.config.poolInfoReady), *scoreBuf)
197+
releaseScoreSlice(scoreBuf)
201198

202199
if obs := observerFromAtomic(&p.observer); obs != nil {
203200
key := indexName + "/" + docID
@@ -219,6 +216,9 @@ func (p *DocRouter) Eval(_ context.Context, req *http.Request) (NextHop, error)
219216

220217
putConnSlice(bp)
221218

219+
if best == nil {
220+
return NextHop{}, nil
221+
}
222222
return NextHop{Conn: best}, nil
223223
}
224224

‎opensearchtransport/policy_index_router.go‎

Lines changed: 75 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,66 @@ var shardCostForWrites = shardCostMultiplier{
8787
shardCostRelocating: costRelocating, // shard moving, may proxy
8888
}
8989

90+
// scoreSliceInitialCap is the starting capacity for pooled score buffers.
91+
// Covers the common case (shard-exact: 1-3 replicas, rendezvous fan-out: 1-3)
92+
// without ratcheting. The cluster-lookup path (all active connections) will
93+
// ratchet up on its first request.
94+
const scoreSliceInitialCap = 4
95+
96+
// scoreSlicePool provides reusable []float64 buffers for [connScoreSelect].
97+
// The target capacity starts at [scoreSliceInitialCap] and ratchets up via
98+
// atomic CAS when larger pools are encountered. On put, slices whose capacity
99+
// doesn't match the current target are discarded so the pool converges on
100+
// the working set size.
101+
//
102+
//nolint:gochecknoglobals // Package-level pool shared by all scoring call sites.
103+
var scoreSlicePool struct {
104+
pool sync.Pool
105+
capacity atomic.Int32
106+
}
107+
108+
//nolint:gochecknoinits // One-time pool initialization.
109+
func init() {
110+
scoreSlicePool.capacity.Store(scoreSliceInitialCap)
111+
scoreSlicePool.pool.New = func() any {
112+
s := make([]float64, 0, int(scoreSlicePool.capacity.Load()))
113+
return &s
114+
}
115+
}
116+
117+
// acquireScoreSlice returns a []float64 of length n from the pool, growing the
118+
// pool's target capacity as needed. The returned pointer must be passed to
119+
// [releaseScoreSlice] when scoring is complete.
120+
func acquireScoreSlice(n int) *[]float64 {
121+
for {
122+
cur := scoreSlicePool.capacity.Load()
123+
if int(cur) >= n {
124+
break
125+
}
126+
if scoreSlicePool.capacity.CompareAndSwap(cur, int32(n)) { //nolint:gosec // n is a connection slice length, bounded well below int32 max
127+
break
128+
}
129+
}
130+
131+
bp := scoreSlicePool.pool.Get().(*[]float64)
132+
if cap(*bp) < n {
133+
*bp = make([]float64, n)
134+
} else {
135+
*bp = (*bp)[:n]
136+
}
137+
return bp
138+
}
139+
140+
// releaseScoreSlice returns a score buffer to the pool. Buffers whose capacity
141+
// doesn't match the current target are discarded (GC'd) so the pool
142+
// converges to a uniform size.
143+
func releaseScoreSlice(bp *[]float64) {
144+
if cap(*bp) != int(scoreSlicePool.capacity.Load()) {
145+
return
146+
}
147+
scoreSlicePool.pool.Put(bp)
148+
}
149+
90150
// forNode returns the shard cost multiplier for a node based on its shard
91151
// composition for the target index.
92152
//
@@ -232,12 +292,9 @@ func (p *IndexRouter) Eval(_ context.Context, req *http.Request) (NextHop, error
232292
return NextHop{}, nil
233293
}
234294

235-
var scoresBuf [8]float64
236-
scores := scoresBuf[:len(conns)]
237-
if len(conns) > len(scoresBuf) {
238-
scores = make([]float64, len(conns))
239-
}
240-
best := connScoreSelect(conns, nil, nil, p.shardCosts, "", loadPoolInfoReady(p.config.poolInfoReady), scores)
295+
scoreBuf := acquireScoreSlice(len(conns))
296+
best := connScoreSelect(conns, nil, nil, p.shardCosts, "", loadPoolInfoReady(p.config.poolInfoReady), *scoreBuf)
297+
releaseScoreSlice(scoreBuf)
241298
if best == nil {
242299
return NextHop{}, nil
243300
}
@@ -260,12 +317,9 @@ func (p *IndexRouter) Eval(_ context.Context, req *http.Request) (NextHop, error
260317
shardCandidates, shardNum, shard := shardExactCandidates(p.cache.features, slot, routingValue, conns)
261318
if len(shardCandidates) > 0 {
262319
// Shard-exact path: score the shard-hosting candidates directly.
263-
var scoresBuf [8]float64
264-
scores := scoresBuf[:len(shardCandidates)]
265-
if len(shardCandidates) > len(scoresBuf) {
266-
scores = make([]float64, len(shardCandidates))
267-
}
268-
best := connScoreSelect(shardCandidates, slot, shard, p.shardCosts, "", loadPoolInfoReady(p.config.poolInfoReady), scores)
320+
scoreBuf := acquireScoreSlice(len(shardCandidates))
321+
best := connScoreSelect(shardCandidates, slot, shard, p.shardCosts, "", loadPoolInfoReady(p.config.poolInfoReady), *scoreBuf)
322+
releaseScoreSlice(scoreBuf)
269323

270324
if obs := observerFromAtomic(&p.observer); obs != nil {
271325
obs.OnRoute(buildRouteEvent(routeEventParams{
@@ -286,6 +340,9 @@ func (p *IndexRouter) Eval(_ context.Context, req *http.Request) (NextHop, error
286340
}))
287341
}
288342

343+
if best == nil {
344+
return NextHop{}, nil
345+
}
289346
return NextHop{Conn: best}, nil
290347
}
291348

@@ -314,12 +371,9 @@ func (p *IndexRouter) Eval(_ context.Context, req *http.Request) (NextHop, error
314371
slot.updateSmoothedMaxBucket(float64(maxBucket))
315372

316373
// Select best candidate with warmup-aware skip/accept.
317-
var scoresBuf [8]float64
318-
scores := scoresBuf[:len(candidates)]
319-
if len(candidates) > len(scoresBuf) {
320-
scores = make([]float64, len(candidates))
321-
}
322-
best := connScoreSelect(candidates, slot, nil, p.shardCosts, "", loadPoolInfoReady(p.config.poolInfoReady), scores)
374+
scoreBuf := acquireScoreSlice(len(candidates))
375+
best := connScoreSelect(candidates, slot, nil, p.shardCosts, "", loadPoolInfoReady(p.config.poolInfoReady), *scoreBuf)
376+
releaseScoreSlice(scoreBuf)
323377

324378
if obs := observerFromAtomic(&p.observer); obs != nil {
325379
obs.OnRoute(buildRouteEvent(routeEventParams{
@@ -340,6 +394,9 @@ func (p *IndexRouter) Eval(_ context.Context, req *http.Request) (NextHop, error
340394

341395
putConnSlice(bp)
342396

397+
if best == nil {
398+
return NextHop{}, nil
399+
}
343400
return NextHop{Conn: best}, nil
344401
}
345402

‎opensearchtransport/policy_index_router_internal_test.go‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
package opensearchtransport
88

99
import (
10+
"fmt"
1011
"net/url"
1112
"testing"
1213
"time"
@@ -217,3 +218,125 @@ func TestCalcConnScore_InFlightChangesScores(t *testing.T) {
217218

218219
require.Less(t, idleScore, busyScore, "idle should score lower (better) than busy")
219220
}
221+
222+
// --- Score slice pool tests ---
223+
224+
func TestAcquireReleaseScoreSlice(t *testing.T) {
225+
t.Parallel()
226+
227+
tests := []struct {
228+
name string
229+
n int
230+
}{
231+
{name: "under initial cap", n: 2},
232+
{name: "at initial cap", n: scoreSliceInitialCap},
233+
{name: "over initial cap", n: scoreSliceInitialCap + 4},
234+
{name: "large pool", n: 64},
235+
}
236+
237+
for _, tt := range tests {
238+
t.Run(tt.name, func(t *testing.T) {
239+
t.Parallel()
240+
bp := acquireScoreSlice(tt.n)
241+
require.NotNil(t, bp)
242+
require.Len(t, *bp, tt.n)
243+
require.GreaterOrEqual(t, cap(*bp), tt.n)
244+
releaseScoreSlice(bp)
245+
})
246+
}
247+
}
248+
249+
func TestScoreSlicePoolRatchetsUp(t *testing.T) {
250+
t.Parallel()
251+
252+
before := scoreSlicePool.capacity.Load()
253+
254+
n := int(before) + 10
255+
bp := acquireScoreSlice(n)
256+
require.Len(t, *bp, n)
257+
258+
after := scoreSlicePool.capacity.Load()
259+
require.GreaterOrEqual(t, int(after), n,
260+
"capacity should ratchet up to at least %d", n)
261+
releaseScoreSlice(bp)
262+
}
263+
264+
func TestScoreSlicePoolDiscardsUndersized(t *testing.T) {
265+
t.Parallel()
266+
267+
// Acquire a small buffer, then ratchet the pool larger.
268+
small := acquireScoreSlice(2)
269+
require.Len(t, *small, 2)
270+
271+
// Force a ratchet by acquiring something larger.
272+
large := acquireScoreSlice(100)
273+
releaseScoreSlice(large)
274+
275+
// Release the small buffer — should be discarded (not panic).
276+
releaseScoreSlice(small)
277+
}
278+
279+
func TestConnScoreSelectLargePool(t *testing.T) {
280+
t.Parallel()
281+
282+
// Regression: the old fixed [8]float64 buffer panicked when
283+
// len(candidates) exceeded 8. Exercise various sizes above that.
284+
tests := []struct {
285+
name string
286+
n int
287+
}{
288+
{name: "9 candidates", n: 9},
289+
{name: "16 candidates", n: 16},
290+
{name: "32 candidates", n: 32},
291+
{name: "64 candidates", n: 64},
292+
}
293+
294+
for _, tt := range tests {
295+
t.Run(tt.name, func(t *testing.T) {
296+
t.Parallel()
297+
298+
conns := make([]*Connection, tt.n)
299+
for i := range tt.n {
300+
conns[i] = scoreTestConn(t, fmt.Sprintf("node-%d", i), 200*time.Microsecond, 0)
301+
conns[i].state.Store(int64(newConnState(lcActive)))
302+
}
303+
304+
scoreBuf := acquireScoreSlice(len(conns))
305+
best := connScoreSelect(conns, nil, nil, &shardCostForReads, "", true, *scoreBuf)
306+
releaseScoreSlice(scoreBuf)
307+
308+
require.NotNil(t, best)
309+
require.Contains(t, conns, best)
310+
})
311+
}
312+
}
313+
314+
// --- Score slice pool benchmarks ---
315+
316+
func BenchmarkAcquireReleaseScoreSlice(b *testing.B) {
317+
tests := []struct {
318+
name string
319+
n int
320+
}{
321+
{name: "4_conns", n: 4},
322+
{name: "8_conns", n: 8},
323+
{name: "16_conns", n: 16},
324+
{name: "64_conns", n: 64},
325+
}
326+
327+
for _, tt := range tests {
328+
b.Run(tt.name, func(b *testing.B) {
329+
// Warm-up: ratchet the pool to the target size so
330+
// steady-state iterations hit the sync.Pool fast path.
331+
warmup := acquireScoreSlice(tt.n)
332+
releaseScoreSlice(warmup)
333+
334+
b.ReportAllocs()
335+
b.ResetTimer()
336+
for range b.N {
337+
bp := acquireScoreSlice(tt.n)
338+
releaseScoreSlice(bp)
339+
}
340+
})
341+
}
342+
}

‎opensearchtransport/policy_pool_router.go‎

Lines changed: 13 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -112,15 +112,12 @@ func (p *poolRouter) Eval(ctx context.Context, req *http.Request) (NextHop, erro
112112
return hop, nil
113113
}
114114

115-
var scoresBuf [8]float64
116-
scores := scoresBuf[:len(conns)]
117-
if len(conns) > len(scoresBuf) {
118-
scores = make([]float64, len(conns))
119-
}
115+
scoreBuf := acquireScoreSlice(len(conns))
120116
pir := loadPoolInfoReady(p.poolInfoReady)
121-
best := connScoreSelect(conns, nil, nil, p.shardCosts, p.poolName, pir, scores)
117+
best := connScoreSelect(conns, nil, nil, p.shardCosts, p.poolName, pir, *scoreBuf)
118+
releaseScoreSlice(scoreBuf)
122119

123-
if best.loadConnState().lifecycle()&(lcActive|lcStandby) == 0 {
120+
if best == nil || best.loadConnState().lifecycle()&(lcActive|lcStandby) == 0 {
124121
return hop, nil
125122
}
126123

@@ -161,13 +158,10 @@ func (p *poolRouter) Eval(ctx context.Context, req *http.Request) (NextHop, erro
161158
effectiveRoutingKey = keyB // OpenSearch default: _id is the routing value
162159
}
163160
shardCandidates, shardNum, shard := shardExactCandidates(p.cache.features, slot, effectiveRoutingKey, conns)
164-
if len(shardCandidates) > 0 { //nolint:nestif // shard-exact path has scoring and observer notification
165-
var scoresBuf [8]float64
166-
scores := scoresBuf[:len(shardCandidates)]
167-
if len(shardCandidates) > len(scoresBuf) {
168-
scores = make([]float64, len(shardCandidates))
169-
}
170-
best := connScoreSelect(shardCandidates, slot, shard, p.shardCosts, p.poolName, loadPoolInfoReady(p.poolInfoReady), scores)
161+
if len(shardCandidates) > 0 {
162+
scoreBuf := acquireScoreSlice(len(shardCandidates))
163+
best := connScoreSelect(shardCandidates, slot, shard, p.shardCosts, p.poolName, loadPoolInfoReady(p.poolInfoReady), *scoreBuf)
164+
releaseScoreSlice(scoreBuf)
171165

172166
if obs := observerFromAtomic(&p.observer); obs != nil {
173167
key := keyA
@@ -194,7 +188,7 @@ func (p *poolRouter) Eval(ctx context.Context, req *http.Request) (NextHop, erro
194188
}
195189

196190
// Verify the selected connection is still active.
197-
if best.loadConnState().lifecycle()&(lcActive|lcStandby) == 0 {
191+
if best == nil || best.loadConnState().lifecycle()&(lcActive|lcStandby) == 0 {
198192
return hop, nil
199193
}
200194

@@ -223,12 +217,9 @@ func (p *poolRouter) Eval(ctx context.Context, req *http.Request) (NextHop, erro
223217
slot.updateSmoothedMaxBucket(float64(maxBucket))
224218

225219
// Select best candidate with warmup-aware skip/accept.
226-
var scoresBuf [8]float64
227-
scores := scoresBuf[:len(candidates)]
228-
if len(candidates) > len(scoresBuf) {
229-
scores = make([]float64, len(candidates))
230-
}
231-
best := connScoreSelect(candidates, slot, nil, p.shardCosts, p.poolName, loadPoolInfoReady(p.poolInfoReady), scores)
220+
scoreBuf := acquireScoreSlice(len(candidates))
221+
best := connScoreSelect(candidates, slot, nil, p.shardCosts, p.poolName, loadPoolInfoReady(p.poolInfoReady), *scoreBuf)
222+
releaseScoreSlice(scoreBuf)
232223

233224
// Compute adaptive max_concurrent_shard_requests for search requests
234225
// routed through a coordinator (non-shard-exact).
@@ -273,7 +264,7 @@ func (p *poolRouter) Eval(ctx context.Context, req *http.Request) (NextHop, erro
273264
// Verify the selected connection is still active (dirty read).
274265
// If it was demoted since the last DiscoveryUpdate, fall through
275266
// to the inner policy's result.
276-
if best.loadConnState().lifecycle()&(lcActive|lcStandby) == 0 {
267+
if best == nil || best.loadConnState().lifecycle()&(lcActive|lcStandby) == 0 {
277268
return hop, nil
278269
}
279270

0 commit comments

Comments
 (0)