1 files changed, 29 insertions, 43 deletions

diff --git a/common/mclock/simclock.go b/common/mclock/simclock.go
index af0f71c43..4d351252f 100644
--- a/common/mclock/simclock.go
+++ b/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,44 +132,31 @@ 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}
-}
-
-func (s *Simulated) init() {
-	if s.cond == nil {
-		s.cond = sync.NewCond(&s.mu)
-	}
+	return ev
 }
 
-// Cancel implements Event.
-func (e *SimulatedEvent) Cancel() bool {
-	s := e.s
+func (ev *simTimer) Stop() bool {
+	s := ev.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
+	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
 		}
 	}
-	if l >= ll || s.scheduled[l].id != e.id {
-		return false
+	return false
+}
+
+func (s *Simulated) init() {
+	if s.cond == nil {
+		s.cond = sync.NewCond(&s.mu)
 	}
-	copy(s.scheduled[l:ll-1], s.scheduled[l+1:])
-	s.scheduled = s.scheduled[:ll-1]
-	return true
 }