Merge pull request #606 from peco/fix-frozen-done-channel

Create MarkComplete() method to avoid future us from deadlocking

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()

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
 

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.

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()

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")
+		}
+	})
+}