diff --git a/action.go b/action.go index 60fbda7..1d6b295 100644 --- a/action.go +++ b/action.go @@ -918,7 +918,7 @@ func doFreezeResults(ctx context.Context, state *Peco, _ Event) { frozen.AppendLine(l) } } - close(frozen.done) + frozen.MarkComplete() state.SetFrozenSource(frozen) state.Query().Reset() diff --git a/action_test.go b/action_test.go index 314a702..e17b41f 100644 --- a/action_test.go +++ b/action_test.go @@ -670,7 +670,7 @@ func TestDoFreezeResults(t *testing.T) { lines := makeLines("frozen1", "frozen2") frozen := NewMemoryBuffer(0) frozen.lines = lines - close(frozen.done) + frozen.MarkComplete() state.SetFrozenSource(frozen) state.ResetCurrentLineBuffer() @@ -695,7 +695,7 @@ func TestDoFreezeResults(t *testing.T) { frozen := NewMemoryBuffer(0) frozen.lines = makeLines("frozen1") - close(frozen.done) + frozen.MarkComplete() state.SetFrozenSource(frozen) state.currentLineBuffer = frozen diff --git a/buffer.go b/buffer.go index 97b2ca6..4d5b27c 100644 --- a/buffer.go +++ b/buffer.go @@ -1,6 +1,7 @@ package peco import ( + "sync" "time" "context" @@ -102,9 +103,22 @@ func (mb *MemoryBuffer) Reset() { defer g.End() } mb.done = make(chan struct{}) + mb.doneOnce = sync.Once{} mb.lines = []line.Line(nil) } +// MarkComplete signals that the buffer is fully populated. It is safe +// to call multiple times; only the first call closes the done channel. +// Use this instead of manually closing the done channel when populating +// a MemoryBuffer outside of the pipeline (e.g. freeze). +func (mb *MemoryBuffer) MarkComplete() { + mb.doneOnce.Do(func() { + mb.mutex.Lock() + close(mb.done) + mb.mutex.Unlock() + }) +} + func (mb *MemoryBuffer) Done() <-chan struct{} { mb.mutex.RLock() defer mb.mutex.RUnlock() @@ -116,11 +130,7 @@ func (mb *MemoryBuffer) Accept(ctx context.Context, in chan interface{}, _ pipel g := pdebug.Marker("MemoryBuffer.Accept") defer g.End() } - defer func() { - mb.mutex.Lock() - close(mb.done) - mb.mutex.Unlock() - }() + defer mb.MarkComplete() // batch collects lines from the channel so we can append them // under a single lock acquisition instead of locking per line. diff --git a/interface.go b/interface.go index e106dc8..8b24656 100644 --- a/interface.go +++ b/interface.go @@ -516,6 +516,7 @@ type Buffer interface { // MemoryBuffer is an implementation of Buffer type MemoryBuffer struct { done chan struct{} + doneOnce sync.Once lines []line.Line mutex sync.RWMutex PeriodicFunc func() diff --git a/peco_test.go b/peco_test.go index e05faa3..24e0a73 100644 --- a/peco_test.go +++ b/peco_test.go @@ -790,3 +790,56 @@ func TestPrintQuery(t *testing.T) { } }) } + +func TestMemoryBufferMarkComplete(t *testing.T) { + t.Run("signals done channel", func(t *testing.T) { + mb := NewMemoryBuffer(0) + mb.MarkComplete() + + select { + case <-mb.Done(): + // expected + default: + t.Fatal("Done channel should be closed after MarkComplete") + } + }) + + t.Run("idempotent - multiple calls do not panic", func(t *testing.T) { + mb := NewMemoryBuffer(0) + mb.MarkComplete() + mb.MarkComplete() // must not panic + mb.MarkComplete() // must not panic + + select { + case <-mb.Done(): + // expected + default: + t.Fatal("Done channel should be closed after MarkComplete") + } + }) + + t.Run("works after Reset", func(t *testing.T) { + mb := NewMemoryBuffer(0) + mb.MarkComplete() + + mb.Reset() + + // After reset, done should be a new open channel + select { + case <-mb.Done(): + t.Fatal("Done channel should not be closed after Reset") + default: + // expected + } + + // MarkComplete should work again on the new channel + mb.MarkComplete() + + select { + case <-mb.Done(): + // expected + default: + t.Fatal("Done channel should be closed after second MarkComplete") + } + }) +}