Skip to content

Commit a9a3c3f

Browse files
committed
fix: optimize WantConnQueue benchmarks to prevent memory exhaustion
- Fix BenchmarkWantConnQueue_Dequeue timeout issue by limiting pre-population - Use object pooling in BenchmarkWantConnQueue_Enqueue to reduce allocations - Optimize BenchmarkWantConnQueue_EnqueueDequeue with reusable wantConn pool - Prevent GitHub Actions benchmark failures due to excessive memory usage Before: BenchmarkWantConnQueue_Dequeue ran for 11+ minutes and was killed After: All benchmarks complete in ~8 seconds with consistent performance
1 parent 56eb8df commit a9a3c3f

File tree

1 file changed

+43
-11
lines changed

1 file changed

+43
-11
lines changed

internal/pool/want_conn_test.go

Lines changed: 43 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -366,22 +366,33 @@ func TestWantConnQueue_ThreadSafety(t *testing.T) {
366366
// Benchmark tests
367367
func BenchmarkWantConnQueue_Enqueue(b *testing.B) {
368368
q := newWantConnQueue()
369+
370+
// Pre-allocate a pool of wantConn to reuse
371+
const poolSize = 1000
372+
wantConnPool := make([]*wantConn, poolSize)
373+
for i := 0; i < poolSize; i++ {
374+
wantConnPool[i] = &wantConn{
375+
ctx: context.Background(),
376+
result: make(chan wantConnResult, 1),
377+
}
378+
}
379+
369380
b.ResetTimer()
370381

371382
for i := 0; i < b.N; i++ {
372-
w := &wantConn{
373-
ctx: context.Background(),
374-
result: make(chan wantConnResult, 1),
375-
}
383+
w := wantConnPool[i%poolSize]
376384
q.enqueue(w)
377385
}
378386
}
379387

380388
func BenchmarkWantConnQueue_Dequeue(b *testing.B) {
381389
q := newWantConnQueue()
382390

383-
// Pre-populate queue
384-
for i := 0; i < b.N; i++ {
391+
// Use a reasonable fixed size for pre-population to avoid memory issues
392+
const queueSize = 10000
393+
394+
// Pre-populate queue with a fixed reasonable size
395+
for i := 0; i < queueSize; i++ {
385396
w := &wantConn{
386397
ctx: context.Background(),
387398
result: make(chan wantConnResult, 1),
@@ -390,20 +401,41 @@ func BenchmarkWantConnQueue_Dequeue(b *testing.B) {
390401
}
391402

392403
b.ResetTimer()
404+
405+
// Benchmark dequeue operations, refilling as needed
393406
for i := 0; i < b.N; i++ {
394-
q.dequeue()
407+
if _, ok := q.dequeue(); !ok {
408+
// Queue is empty, refill a batch
409+
for j := 0; j < 1000; j++ {
410+
w := &wantConn{
411+
ctx: context.Background(),
412+
result: make(chan wantConnResult, 1),
413+
}
414+
q.enqueue(w)
415+
}
416+
// Dequeue again
417+
q.dequeue()
418+
}
395419
}
396420
}
397421

398422
func BenchmarkWantConnQueue_EnqueueDequeue(b *testing.B) {
399423
q := newWantConnQueue()
400-
b.ResetTimer()
401-
402-
for i := 0; i < b.N; i++ {
403-
w := &wantConn{
424+
425+
// Pre-allocate a pool of wantConn to reuse
426+
const poolSize = 1000
427+
wantConnPool := make([]*wantConn, poolSize)
428+
for i := 0; i < poolSize; i++ {
429+
wantConnPool[i] = &wantConn{
404430
ctx: context.Background(),
405431
result: make(chan wantConnResult, 1),
406432
}
433+
}
434+
435+
b.ResetTimer()
436+
437+
for i := 0; i < b.N; i++ {
438+
w := wantConnPool[i%poolSize]
407439
q.enqueue(w)
408440
q.dequeue()
409441
}

0 commit comments

Comments
 (0)