Merge pull request #4 from vimeo/fakeclock_abort_sleep_counters

fake clock: Add tracking for sleep/abort-counters

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,24 +170,30 @@ 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()
 }
 
 // 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