Add SleepNotify, to notify a channel once the "waiter" is fully registered but before actually waiting for the Sleep to complete

Author: alaingilbert

clockwork.go

@@ -121,7 +121,15 @@ func (fc *FakeClock) After(d time.Duration) <-chan time.Time {
 
 // Sleep blocks until the given duration has passed on the fakeClock.
 func (fc *FakeClock) Sleep(d time.Duration) {
-	<-fc.After(d)
+	fc.SleepNotify(d, make(chan struct{}))
+}
+
+// SleepNotify blocks until the given duration has passed on the fakeClock.
+// Notify "ch" once the waiters has been updated
+func (fc *FakeClock) SleepNotify(d time.Duration, ch chan struct{}) {
+	done := fc.After(d)
+	close(ch)
+	<-done
 }
 
 // Now returns the current time of the fakeClock

clockwork_test.go

@@ -3,6 +3,8 @@ package clockwork
 import (
 	"context"
 	"errors"
+	"math/rand"
+	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -209,3 +211,26 @@ func TestFakeClockRace(t *testing.T) {
 	go func() { fc.NewTimer(d) }()
 	go func() { fc.Sleep(d) }()
 }
+
+func TestSleepNotify(t *testing.T) {
+	var calls atomic.Int32
+	clock := NewFakeClock()
+	beforeCh := make(chan struct{})
+	afterCh := make(chan struct{})
+	go func() { // thread #1
+		clock.SleepNotify(time.Minute, beforeCh) // We want to wait for this before advancing the clock
+		calls.Add(1)
+		close(afterCh)
+	}()
+	go func() {                // thread #2
+		if rand.Intn(2) == 0 { // 50% chance of making another Sleep
+			clock.Sleep(time.Hour)
+		}
+	}()
+	<-beforeCh
+	clock.Advance(time.Minute)
+	<-afterCh
+	if calls.Load() != 1 {
+		t.Fatalf("SleepNotify() did not call the callback")
+	}
+}