common/mclock: clean up AfterFunc support (#20054)

This change adds tests for the virtual clock and aligns the interface
with the time package by renaming Cancel to Stop. It also removes the
binary search from Stop because it complicates the code unnecessarily.

common/mclock/simclock.go

@@ -32,22 +32,17 @@ import (
 // the timeout using a channel or semaphore.
 type Simulated struct {
 	now       AbsTime
-	scheduled []event
+	scheduled []*simTimer
 	mu        sync.RWMutex
 	cond      *sync.Cond
 	lastId    uint64
 }
 
-type event struct {
+// simTimer implements Timer on the virtual clock.
+type simTimer struct {
 	do func()
 	at AbsTime
 	id uint64
-}
-
-// SimulatedEvent implements Event for a virtual clock.
-type SimulatedEvent struct {
-	at AbsTime
-	id uint64
 	s  *Simulated
 }
 
@@ -75,6 +70,7 @@ func (s *Simulated) Run(d time.Duration) {
 	}
 }
 
+// ActiveTimers returns the number of timers that haven't fired.
 func (s *Simulated) ActiveTimers() int {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
@@ -82,6 +78,7 @@ func (s *Simulated) ActiveTimers() int {
 	return len(s.scheduled)
 }
 
+// WaitForTimers waits until the clock has at least n scheduled timers.
 func (s *Simulated) WaitForTimers(n int) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -92,7 +89,7 @@ func (s *Simulated) WaitForTimers(n int) {
 	}
 }
 
-// Now implements Clock.
+// Now returns the current virtual time.
 func (s *Simulated) Now() AbsTime {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
@@ -100,12 +97,13 @@ func (s *Simulated) Now() AbsTime {
 	return s.now
 }
 
-// Sleep implements Clock.
+// Sleep blocks until the clock has advanced by d.
 func (s *Simulated) Sleep(d time.Duration) {
 	<-s.After(d)
 }
 
-// After implements Clock.
+// After returns a channel which receives the current time after the clock
+// has advanced by d.
 func (s *Simulated) After(d time.Duration) <-chan time.Time {
 	after := make(chan time.Time, 1)
 	s.AfterFunc(d, func() {
@@ -114,8 +112,9 @@ func (s *Simulated) After(d time.Duration) <-chan time.Time {
 	return after
 }
 
-// AfterFunc implements Clock.
-func (s *Simulated) AfterFunc(d time.Duration, do func()) Event {
+// AfterFunc runs fn after the clock has advanced by d. Unlike with the system
+// clock, fn runs on the goroutine that calls Run.
+func (s *Simulated) AfterFunc(d time.Duration, fn func()) Timer {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	s.init()
@@ -133,12 +132,27 @@ func (s *Simulated) AfterFunc(d time.Duration, do func()) Event {
 			l = m + 1
 		}
 	}
-	s.scheduled = append(s.scheduled, event{})
+	ev := &simTimer{do: fn, at: at, s: s}
+	s.scheduled = append(s.scheduled, nil)
 	copy(s.scheduled[l+1:], s.scheduled[l:ll])
-	e := event{do: do, at: at, id: id}
-	s.scheduled[l] = e
+	s.scheduled[l] = ev
 	s.cond.Broadcast()
-	return &SimulatedEvent{at: at, id: id, s: s}
+	return ev
+}
+
+func (ev *simTimer) Stop() bool {
+	s := ev.s
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	for i := 0; i < len(s.scheduled); i++ {
+		if s.scheduled[i] == ev {
+			s.scheduled = append(s.scheduled[:i], s.scheduled[i+1:]...)
+			s.cond.Broadcast()
+			return true
+		}
+	}
+	return false
 }
 
 func (s *Simulated) init() {
@@ -146,31 +160,3 @@ func (s *Simulated) init() {
 		s.cond = sync.NewCond(&s.mu)
 	}
 }
-
-// Cancel implements Event.
-func (e *SimulatedEvent) Cancel() bool {
-	s := e.s
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
-	l, h := 0, len(s.scheduled)
-	ll := h
-	for l != h {
-		m := (l + h) / 2
-		if e.id == s.scheduled[m].id {
-			l = m
-			break
-		}
-		if (e.at < s.scheduled[m].at) || ((e.at == s.scheduled[m].at) && (e.id < s.scheduled[m].id)) {
-			h = m
-		} else {
-			l = m + 1
-		}
-	}
-	if l >= ll || s.scheduled[l].id != e.id {
-		return false
-	}
-	copy(s.scheduled[l:ll-1], s.scheduled[l+1:])
-	s.scheduled = s.scheduled[:ll-1]
-	return true
-}