Add a WithSynchronousAfterFunc option

clockwork.go

@@ -41,20 +41,36 @@ func NewRealClock() Clock {
 	return &realClock{}
 }
 
+// FakeClockOption is an option for configuring a FakeClock.
+type FakeClockOption func(*fakeClock)
+
+// WithSynchronousAfterFunc is a FakeClockOption which ensures that the callback to AfterFunc is called synchronously
+// instead of in a separate goroutine. If true, this option will cause any Advance calls to block until the callback
+// from all fired timers has completed.
+func WithSynchronousAfterFunc(b bool) FakeClockOption {
+	return func(c *fakeClock) {
+		c.synchronousAfterFunc = b
+	}
+}
+
 // NewFakeClock returns a FakeClock implementation which can be
 // manually advanced through time for testing. The initial time of the
 // FakeClock will be the current system time.
 //
 // Tests that require a deterministic time must use NewFakeClockAt.
-func NewFakeClock() FakeClock {
-	return NewFakeClockAt(time.Now())
+func NewFakeClock(opts ...FakeClockOption) FakeClock {
+	return NewFakeClockAt(time.Now(), opts...)
 }
 
 // NewFakeClockAt returns a FakeClock initialised at the given time.Time.
-func NewFakeClockAt(t time.Time) FakeClock {
-	return &fakeClock{
+func NewFakeClockAt(t time.Time, opts ...FakeClockOption) FakeClock {
+	clock := &fakeClock{
 		time: t,
 	}
+	for _, opt := range opts {
+		opt(clock)
+	}
+	return clock
 }
 
 type realClock struct{}
@@ -94,6 +110,8 @@ type fakeClock struct {
 	waiters  []expirer
 	blockers []*blocker
 	time     time.Time
+	// synchronousAfterFunc determines whether AfterFunc calls should be executed synchronously or in a goroutine.
+	synchronousAfterFunc bool
 }
 
 // blocker is a caller of BlockUntil.
@@ -181,7 +199,8 @@ func (fc *fakeClock) newTimer(d time.Duration, afterfunc func()) *fakeTimer {
 		},
 		stop: func() bool { return fc.stop(ft) },
 
-		afterFunc: afterfunc,
+		afterFunc:            afterfunc,
+		synchronousAfterFunc: fc.synchronousAfterFunc,
 	}
 	fc.set(ft, d)
 	return ft
@@ -307,7 +326,7 @@ func (fc *fakeClock) setExpirer(e expirer, d time.Duration) {
 		return fc.waiters[i].expiry().Before(fc.waiters[j].expiry())
 	})
 
-    // Notify blockers of our new waiter.
+	// Notify blockers of our new waiter.
 	var blocked []*blocker
 	count := len(fc.waiters)
 	for _, b := range fc.blockers {

clockwork_test.go

@@ -3,7 +3,9 @@ package clockwork
 import (
 	"context"
 	"errors"
+	"fmt"
 	"reflect"
+	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -199,3 +201,33 @@ func TestFakeClockRace(t *testing.T) {
 	go func() { fc.NewTimer(d) }()
 	go func() { fc.Sleep(d) }()
 }
+
+func TestFakeClockSynchronousAfterFunc(t *testing.T) {
+	t.Parallel()
+	fc := NewFakeClock(WithSynchronousAfterFunc(true))
+	var called atomic.Bool
+	fc.AfterFunc(time.Second, func() {
+		called.Store(true)
+	})
+	if called.Load() {
+		t.Fatal("AfterFunc called before Advance")
+	}
+	fc.Advance(time.Second)
+	if !called.Load() {
+		t.Fatal("AfterFunc not called after Advance")
+	}
+}
+
+func ExampleWithSynchronousAfterFunc() {
+	fc := NewFakeClock(WithSynchronousAfterFunc(true))
+	fc.AfterFunc(time.Second, func() {
+		fmt.Println("AfterFunc called")
+	})
+	fmt.Println("Calling Advance")
+	fc.Advance(time.Second)
+	fmt.Println("Advance returned")
+	// Output:
+	// Calling Advance
+	// AfterFunc called
+	// Advance returned
+}

example_test.go

@@ -46,3 +46,29 @@ func TestMyFunc(t *testing.T) {
 	// Assert the final state.
 	assertState(t, i, 1)
 }
+
+// myFunc2 is an example of a time-dependent function which uses AfterFunc.
+func myFunc2(clock Clock, i *int) {
+	clock.AfterFunc(3*time.Second, func() {
+		*i += 1
+	})
+}
+
+func TestMyFunc2(t *testing.T) {
+	var i int
+
+	// Use WithSynchronousAfterFunc to ensure that the AfterFunc callback is called by the time Advance returns.
+	c := NewFakeClock(WithSynchronousAfterFunc(true))
+
+	// Call myFunc2, which will schedule a callback to increment i.
+	myFunc2(c, &i)
+
+	// Assert the initial state.
+	assertState(t, i, 0)
+
+	// Now advance the clock forward in time to trigger the callback.
+	c.Advance(3 * time.Second)
+
+	// Assert the final state.
+	assertState(t, i, 1)
+}

timer.go

@@ -28,6 +28,8 @@ type fakeTimer struct {
 	// If present when the timer fires, the timer calls afterFunc in its own
 	// goroutine rather than sending the time on Chan().
 	afterFunc func()
+	// synchronousAfterFunc indicates that the afterFunc should be called in the same goroutine as the timer.
+	synchronousAfterFunc bool
 }
 
 func (f *fakeTimer) Reset(d time.Duration) bool {
@@ -40,7 +42,11 @@ func (f *fakeTimer) Stop() bool {
 
 func (f *fakeTimer) expire(now time.Time) *time.Duration {
 	if f.afterFunc != nil {
-		go f.afterFunc()
+		if f.synchronousAfterFunc {
+			f.afterFunc()
+		} else {
+			go f.afterFunc()
+		}
 		return nil
 	}