Merge pull request #690 from peco/fix-hub-batch-locking

Replace shouldLock with context-based re-entrancy

action.go

@@ -422,14 +422,10 @@ func doCancel(ctx context.Context, state *Peco, e Event) {
 	state.Exit(err)
 }
 
-// batchAction extracts the top-level action call flag from the context,
-// runs fn inside a Hub.Batch, and marks nested calls as non-top-level.
+// batchAction runs fn inside a Hub.Batch. Re-entrant locking is
+// handled automatically by Hub.Batch via context detection.
 func batchAction(ctx context.Context, state *Peco, fn func(context.Context)) {
-	toplevel, _ := ctx.Value(isTopLevelActionCallKey{}).(bool)
-	state.Hub().Batch(ctx, func(ctx context.Context) {
-		ctx = context.WithValue(ctx, isTopLevelActionCallKey{}, false)
-		fn(ctx)
-	}, toplevel)
+	state.Hub().Batch(ctx, fn)
 }
 
 func doToggleSelectionAndSelectNext(ctx context.Context, state *Peco, e Event) {

hub/hub.go

@@ -51,16 +51,22 @@ 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{}
+
 // Batch allows you to synchronously send messages during the
-// scope of f() being executed.
-func (h *Hub) Batch(ctx context.Context, f func(ctx context.Context), shouldLock bool) {
+// 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)
+
 	if pdebug.Enabled {
-		g := pdebug.Marker("Batch (shouldLock=%t)", shouldLock)
+		g := pdebug.Marker("Batch (nested=%t)", nested)
 		defer g.End()
 	}
 
-	if shouldLock {
-		// lock during this operation
+	if !nested {
 		h.mutex.Lock()
 		defer h.mutex.Unlock()
 	}
@@ -75,7 +81,9 @@ func (h *Hub) Batch(ctx context.Context, f func(ctx context.Context), shouldLock
 		}
 	}()
 
-	f(context.WithValue(ctx, batchPayloadKey{}, true))
+	batchCtx := context.WithValue(ctx, batchPayloadKey{}, true)
+	batchCtx = context.WithValue(batchCtx, batchLockKey{}, true)
+	f(batchCtx)
 }
 
 var doneChPool = sync.Pool{

hub/hub_test.go

@@ -62,7 +62,7 @@ func TestHub(t *testing.T) {
 		h.SendDraw(ctx, &hub.DrawOptions{})
 		h.SendStatusMsg(ctx, "Hello, World!", 0)
 		h.SendPaging(ctx, hub.PagingRequestType(1))
-	}, true)
+	})
 
 	phases := []string{
 		"query",
@@ -96,7 +96,7 @@ func TestBatchPanicPropagates(t *testing.T) {
 	require.Panics(t, func() {
 		h.Batch(ctx, func(_ context.Context) {
 			panic("bug in callback")
-		}, true)
+		})
 	}, "Batch must not silently swallow panics")
 }
 
@@ -109,7 +109,7 @@ func TestBatchPanicReleasesLock(t *testing.T) {
 		defer func() { recover() }()
 		h.Batch(ctx, func(_ context.Context) {
 			panic("first call panics")
-		}, true)
+		})
 	}()
 
 	// Second call: if the mutex was not released by the first panic,
@@ -123,7 +123,7 @@ func TestBatchPanicReleasesLock(t *testing.T) {
 				p.Done()
 			}()
 			h.SendQuery(ctx, "after panic")
-		}, true)
+		})
 		close(done)
 	}()
 
@@ -135,6 +135,33 @@ func TestBatchPanicReleasesLock(t *testing.T) {
 	}
 }
 
+func TestBatchNestedDoesNotDeadlock(t *testing.T) {
+	h := hub.New(5)
+	ctx := context.Background()
+
+	done := make(chan struct{})
+	go func() {
+		h.Batch(ctx, func(ctx context.Context) {
+			// Nested Batch call — must not deadlock
+			h.Batch(ctx, func(ctx context.Context) {
+				go func() {
+					p := <-h.QueryCh()
+					p.Done()
+				}()
+				h.SendQuery(ctx, "nested")
+			})
+		})
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		// success — nested Batch did not deadlock
+	case <-time.After(2 * time.Second):
+		t.Fatal("nested Batch deadlocked")
+	}
+}
+
 func TestSendStatusMsg(t *testing.T) {
 	t.Run("zero delay", func(t *testing.T) {
 		h := hub.New(5)

interface.go

@@ -494,7 +494,7 @@ type Input struct {
 // Most code (actions, input handling, source setup) only needs
 // the sender side.
 type HubSender interface {
-	Batch(context.Context, func(context.Context), bool)
+	Batch(context.Context, func(context.Context))
 	SendDraw(context.Context, *hub.DrawOptions)
 	SendDrawPrompt(context.Context)
 	SendPaging(context.Context, hub.PagingRequest)

keymap.go

@@ -26,8 +26,6 @@ func (km Keymap) Sequence() Keyseq {
 	return km.seq
 }
 
-type isTopLevelActionCallKey struct{}
-
 func (km Keymap) ExecuteAction(ctx context.Context, state *Peco, ev Event) (err error) {
 	if pdebug.Enabled {
 		g := pdebug.Marker("Keymap.ExecuteAction %v", ev).BindError(&err)
@@ -39,7 +37,6 @@ func (km Keymap) ExecuteAction(ctx context.Context, state *Peco, ev Event) (err
 		return errors.New("action not found")
 	}
 
-	ctx = context.WithValue(ctx, isTopLevelActionCallKey{}, true)
 	a.Execute(ctx, state, ev)
 	return nil
 }

peco.go

@@ -801,7 +801,7 @@ func (p *Peco) sendQuery(ctx context.Context, q string, nextFunc func()) {
 			if nextFunc != nil {
 				nextFunc()
 			}
-		}, false)
+		})
 	}
 }
 
@@ -841,7 +841,7 @@ func (p *Peco) ExecQuery(ctx context.Context, nextFunc func()) bool {
 			if nextFunc != nil {
 				nextFunc()
 			}
-		}, false)
+		})
 		return true
 	}
 

peco_test.go

@@ -24,7 +24,7 @@ import (
 
 type nullHub struct{}
 
-func (h nullHub) Batch(_ context.Context, _ func(context.Context), _ bool)   {}
+func (h nullHub) Batch(_ context.Context, _ func(context.Context))           {}
 func (h nullHub) DrawCh() chan *hub.Payload[*hub.DrawOptions]                { return nil }
 func (h nullHub) PagingCh() chan *hub.Payload[hub.PagingRequest]             { return nil }
 func (h nullHub) QueryCh() chan *hub.Payload[string]                         { return nil }