From e367bbc4eff1a4a80694f909322f29f46a13bd44 Mon Sep 17 00:00:00 2001
From: David Finkel <davidf@vimeo.com>
Date: Mon, 20 Jul 2020 10:44:14 -0400
Subject: [PATCH] fake clock: Add tracking for sleep/abort-counters

It's useful to know how many calls to `SleepUntil` and `SleepFor` have
awoken due to canceled contexts, and sometimes synchronizing tests
requires waiting until some number of goroutines have both awakened and
gone back to sleep. Add primitives to allow this.
---
 clocks/fake_clock.go | 60 ++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 55 insertions(+), 5 deletions(-)

diff --git a/clocks/fake_clock.go b/clocks/fake_clock.go
index d3ed67b..420bc3f 100644
--- a/clocks/fake_clock.go
+++ b/clocks/fake_clock.go
@@ -22,6 +22,13 @@ type FakeClock struct {
 
 	// counter tracking the number of wakeups (protected by mu)
 	wakeups int
+
+	// counter tracking the number of cancelled sleeps (protected by mu)
+	sleepAborts int
+
+	// counter tracking the number of sleepers who have ever gone to sleep
+	// (protected by mu)
+	sleepersAggregate int
 }
 
 // NewFakeClock returns an initialized FakeClock instance.
@@ -75,6 +82,22 @@ func (f *FakeClock) NumSleepers() int {
 	return len(f.sleepers)
 }
 
+// NumAggSleepers returns the number of goroutines who have ever slept under
+// SleepFor and SleepUntil calls.
+func (f *FakeClock) NumAggSleepers() int {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.sleepersAggregate
+}
+
+// NumSleepAborts returns the number of calls to SleepFor and SleepUntil which
+// have ended prematurely due to canceled contexts.
+func (f *FakeClock) NumSleepAborts() int {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.sleepAborts
+}
+
 // Sleepers returns the number of goroutines waiting in SleepFor and SleepUntil
 // calls.
 func (f *FakeClock) Sleepers() []time.Time {
@@ -96,6 +119,26 @@ func (f *FakeClock) AwaitSleepers(n int) {
 	}
 }
 
+// AwaitAggSleepers waits until the aggregate number of sleepers exceeds its
+// argument
+func (f *FakeClock) AwaitAggSleepers(n int) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	for f.sleepersAggregate < n {
+		f.cond.Wait()
+	}
+}
+
+// AwaitSleepAborts waits until the number of aborted sleepers exceeds its
+// argument
+func (f *FakeClock) AwaitSleepAborts(n int) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	for f.sleepAborts < n {
+		f.cond.Wait()
+	}
+}
+
 // Wakeups returns the number of sleepers that have been awoken (useful for
 // verifying that nothing was woken up when advancing time)
 func (f *FakeClock) Wakeups() int {
@@ -127,14 +170,20 @@ func (f *FakeClock) setAbsoluteWaiter(until time.Time) chan struct{} {
 		return ch
 	}
 	f.sleepers[ch] = until
+	f.sleepersAggregate++
 
 	f.cond.Broadcast()
 	return ch
 }
 
-func (f *FakeClock) removeWaiter(ch chan struct{}) {
+func (f *FakeClock) removeWaiter(ch chan struct{}, abort bool) {
 	f.mu.Lock()
 	defer f.mu.Unlock()
+	// If the channel is present, and this was an abort, increment the
+	// aborts counter.
+	if _, ok := f.sleepers[ch]; ok && abort {
+		f.sleepAborts++
+	}
 	delete(f.sleepers, ch)
 	f.cond.Broadcast()
 }
@@ -142,9 +191,9 @@ func (f *FakeClock) removeWaiter(ch chan struct{}) {
 // SleepUntil blocks until either ctx expires or until arrives.
 // Return value is false if context-cancellation/expiry prompted an
 // early return
-func (f *FakeClock) SleepUntil(ctx context.Context, until time.Time) bool {
+func (f *FakeClock) SleepUntil(ctx context.Context, until time.Time) (success bool) {
 	ch := f.setAbsoluteWaiter(until)
-	defer f.removeWaiter(ch)
+	defer func() { f.removeWaiter(ch, !success) }()
 	select {
 	case <-ch:
 		return true
@@ -158,17 +207,18 @@ func (f *FakeClock) setRelativeWaiter(dur time.Duration) chan struct{} {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 	f.sleepers[ch] = f.current.Add(dur)
+	f.sleepersAggregate++
 	f.cond.Broadcast()
 	return ch
 }
 
 // SleepFor is the relative-time equivalent of SleepUntil.
-func (f *FakeClock) SleepFor(ctx context.Context, dur time.Duration) bool {
+func (f *FakeClock) SleepFor(ctx context.Context, dur time.Duration) (success bool) {
 	if dur <= 0 {
 		return true
 	}
 	ch := f.setRelativeWaiter(dur)
-	defer f.removeWaiter(ch)
+	defer func() { f.removeWaiter(ch, !success) }()
 	select {
 	case <-ch:
 		return true