From e2bc390c68590e8b6c4b0639dbbc583a167a8ca0 Mon Sep 17 00:00:00 2001
From: Daisuke Maki <lestrrat+github@gmail.com>
Date: Sat, 21 Feb 2026 22:15:03 +0900
Subject: [PATCH] fix: merge two context keys into one in Hub.Batch

Reduce context.WithValue calls from 2 to 1 per Batch invocation.

| Metric     | Before | After | Delta |
|------------|--------|-------|-------|
| allocs/op  | 3      | 2     | -1    |
| B/op       | 128    | 80    | -48   |
| ns/op      | ~372   | ~352  | -5%   |
---
 hub/bench_test.go | 29 +++++++++++++++++++++++++++++
 hub/hub.go        | 22 ++++++++--------------
 2 files changed, 37 insertions(+), 14 deletions(-)
 create mode 100644 hub/bench_test.go

diff --git a/hub/bench_test.go b/hub/bench_test.go
new file mode 100644
index 0000000..01478bb
--- /dev/null
+++ b/hub/bench_test.go
@@ -0,0 +1,29 @@
+package hub_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/peco/peco/hub"
+)
+
+// BenchmarkHubBatch measures the allocation cost of Hub.Batch context setup.
+func BenchmarkHubBatch(b *testing.B) {
+	h := hub.New(5)
+	ctx := context.Background()
+
+	// Drain query channel and call Done() so batch sends unblock
+	go func() {
+		for p := range h.QueryCh() {
+			p.Done()
+		}
+	}()
+
+	b.ResetTimer()
+	b.ReportAllocs()
+	for b.Loop() {
+		h.Batch(ctx, func(bctx context.Context) {
+			h.SendQuery(bctx, "test")
+		})
+	}
+}
diff --git a/hub/hub.go b/hub/hub.go
index 87eab4d..ddf982a 100644
--- a/hub/hub.go
+++ b/hub/hub.go
@@ -70,17 +70,16 @@ func New(bufsiz int) *Hub {
 	}
 }
 
-type batchPayloadKey struct{}
-
-// batchLockKey is used to detect re-entrant Batch calls so that
-// nested calls skip mutex acquisition and avoid deadlock.
-type batchLockKey struct{}
+// batchCtxKey is a single context key that signals both "this is a batch
+// payload" and "the hub mutex is already held" for re-entrant detection.
+// Using one key instead of two avoids a second context.WithValue allocation.
+type batchCtxKey struct{}
 
 // Batch allows you to synchronously send messages during the
 // scope of f() being executed. The mutex is acquired automatically
 // unless this is a nested Batch call (detected via context).
 func (h *Hub) Batch(ctx context.Context, f func(ctx context.Context)) {
-	nested, _ := ctx.Value(batchLockKey{}).(bool)
+	nested, _ := ctx.Value(batchCtxKey{}).(bool)
 
 	if pdebug.Enabled {
 		g := pdebug.Marker("Batch (nested=%t)", nested)
@@ -102,8 +101,7 @@ func (h *Hub) Batch(ctx context.Context, f func(ctx context.Context)) {
 		}
 	}()
 
-	batchCtx := context.WithValue(ctx, batchPayloadKey{}, true)
-	batchCtx = context.WithValue(batchCtx, batchLockKey{}, true)
+	batchCtx := context.WithValue(ctx, batchCtxKey{}, true)
 	f(batchCtx)
 }
 
@@ -132,12 +130,8 @@ func (p *Payload[T]) waitDone() {
 
 // isBatchCtx reports whether the context was created by a Batch call.
 func isBatchCtx(ctx context.Context) bool {
-	var isBatchMode bool
-	v := ctx.Value(batchPayloadKey{})
-	if vv, ok := v.(bool); ok {
-		isBatchMode = vv
-	}
-	return isBatchMode
+	v, _ := ctx.Value(batchCtxKey{}).(bool)
+	return v
 }
 
 // send is the low-level generic utility for sending typed payloads.