From 556fee8b4e70a645dbede2d83b7ac21a20ca4eda Mon Sep 17 00:00:00 2001
From: Nick Santos <nick@tilt.dev>
Date: Wed, 23 Mar 2022 16:48:05 -0400
Subject: [PATCH] fix: fixes a deadlock in BlockUntil

fixes https://github.com/jonboulle/clockwork/issues/35
---
 clockwork.go      | 14 +++++++-------
 clockwork_test.go | 32 ++++++++++++++++++++++++++++----
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/clockwork.go b/clockwork.go
index 1018051..8c0e257 100644
--- a/clockwork.go
+++ b/clockwork.go
@@ -113,12 +113,12 @@ func (fc *fakeClock) After(d time.Duration) <-chan time.Time {
 	return done
 }
 
-// notifyBlockers notifies all the blockers waiting until the
-// given number of sleepers are waiting on the fakeClock. It
-// returns an updated slice of blockers (i.e. those still waiting)
+// notifyBlockers notifies all the blockers waiting until the at least the given
+// number of sleepers are waiting on the fakeClock. It returns an updated slice
+// of blockers (i.e. those still waiting)
 func notifyBlockers(blockers []*blocker, count int) (newBlockers []*blocker) {
 	for _, b := range blockers {
-		if b.count == count {
+		if b.count <= count {
 			close(b.ch)
 		} else {
 			newBlockers = append(newBlockers, b)
@@ -179,12 +179,12 @@ func (fc *fakeClock) Advance(d time.Duration) {
 // (callers of Sleep or After)
 func (fc *fakeClock) BlockUntil(n int) {
 	fc.l.Lock()
-	// Fast path: current number of sleepers is what we're looking for
-	if len(fc.sleepers) == n {
+	// Fast path: we already have >= n sleepers.
+	if len(fc.sleepers) >= n {
 		fc.l.Unlock()
 		return
 	}
-	// Otherwise, set up a new blocker
+	// Otherwise, we have < n sleepers. Set up a new blocker to wait for more.
 	b := &blocker{
 		count: n,
 		ch:    make(chan struct{}),
diff --git a/clockwork_test.go b/clockwork_test.go
index 6b8b5cf..e5dcf72 100644
--- a/clockwork_test.go
+++ b/clockwork_test.go
@@ -89,8 +89,13 @@ func TestNotifyBlockers(t *testing.T) {
 	b5 := &blocker{10, make(chan struct{})}
 	bs := []*blocker{b1, b2, b3, b4, b5}
 	bs1 := notifyBlockers(bs, 2)
-	if n := len(bs1); n != 4 {
-		t.Fatalf("got %d blockers, want %d", n, 4)
+	if n := len(bs1); n != 3 {
+		t.Fatalf("got %d blockers, want %d", n, 3)
+	}
+	select {
+	case <-b1.ch:
+	case <-time.After(time.Second):
+		t.Fatalf("timed out waiting for channel close!")
 	}
 	select {
 	case <-b2.ch:
@@ -98,8 +103,13 @@ func TestNotifyBlockers(t *testing.T) {
 		t.Fatalf("timed out waiting for channel close!")
 	}
 	bs2 := notifyBlockers(bs1, 10)
-	if n := len(bs2); n != 2 {
-		t.Fatalf("got %d blockers, want %d", n, 2)
+	if n := len(bs2); n != 0 {
+		t.Fatalf("got %d blockers, want %d", n, 0)
+	}
+	select {
+	case <-b3.ch:
+	case <-time.After(time.Second):
+		t.Fatalf("timed out waiting for channel close!")
 	}
 	select {
 	case <-b4.ch:
@@ -144,3 +154,17 @@ func TestFakeClockSince(t *testing.T) {
 		t.Fatalf("fakeClock.Since() returned unexpected duration, got: %d, want: %d", fc.Since(now), elapsedTime)
 	}
 }
+
+// This used to result in a deadlock.
+// https://github.com/jonboulle/clockwork/issues/35
+func TestTwoBlockersOneBlock(t *testing.T) {
+	fc := &fakeClock{}
+
+	ft1 := fc.NewTicker(time.Second)
+	ft2 := fc.NewTicker(time.Second)
+
+	fc.BlockUntil(1)
+	fc.BlockUntil(2)
+	ft1.Stop()
+	ft2.Stop()
+}