Avoid race in diode.Close with waiter (#481)

rs · Sep 15, 2022 · 315967f · 315967f
1 parent e218d18
commit 315967f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/diode/diode_test.go b/diode/diode_test.go
@@ -30,6 +30,14 @@ func TestNewWriter(t *testing.T) {
 	}
 }
 
+func TestClose(t *testing.T) {
+	buf := bytes.Buffer{}
+	w := diode.NewWriter(&buf, 1000, 0, func(missed int) {})
+	log := zerolog.New(w)
+	log.Print("test")
+	w.Close()
+}
+
 func Benchmark(b *testing.B) {
 	log.SetOutput(ioutil.Discard)
 	defer log.SetOutput(os.Stderr)

diff --git a/diode/internal/diodes/waiter.go b/diode/internal/diodes/waiter.go
@@ -39,7 +39,12 @@ func NewWaiter(d Diode, opts ...WaiterConfigOption) *Waiter {
 
 	go func() {
 		<-w.ctx.Done()
+
+		// Mutex is strictly necessary here to avoid a race in Next() (between
+		// w.isDone() and w.c.Wait()) and w.c.Broadcast() here.
+		w.mu.Lock()
 		w.c.Broadcast()
+		w.mu.Unlock()
 	}()
 
 	return w