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.

3 files changed, 155 insertions, 62 deletions

diff --git a/common/mclock/mclock.go b/common/mclock/mclock.go
index 0c941082f..d0e0cd78b 100644
--- a/common/mclock/mclock.go
+++ b/common/mclock/mclock.go
@@ -36,47 +36,39 @@ func (t AbsTime) Add(d time.Duration) AbsTime {
 	return t + AbsTime(d)
 }
 
-// Clock interface makes it possible to replace the monotonic system clock with
+// The Clock interface makes it possible to replace the monotonic system clock with
 // a simulated clock.
 type Clock interface {
 	Now() AbsTime
 	Sleep(time.Duration)
 	After(time.Duration) <-chan time.Time
-	AfterFunc(d time.Duration, f func()) Event
+	AfterFunc(d time.Duration, f func()) Timer
 }
 
-// Event represents a cancellable event returned by AfterFunc
-type Event interface {
-	Cancel() bool
+// Timer represents a cancellable event returned by AfterFunc
+type Timer interface {
+	Stop() bool
 }
 
 // System implements Clock using the system clock.
 type System struct{}
 
-// Now implements Clock.
+// Now returns the current monotonic time.
 func (System) Now() AbsTime {
 	return AbsTime(monotime.Now())
 }
 
-// Sleep implements Clock.
+// Sleep blocks for the given duration.
 func (System) Sleep(d time.Duration) {
 	time.Sleep(d)
 }
 
-// After implements Clock.
+// After returns a channel which receives the current time after d has elapsed.
 func (System) After(d time.Duration) <-chan time.Time {
 	return time.After(d)
 }
 
-// AfterFunc implements Clock.
-func (System) AfterFunc(d time.Duration, f func()) Event {
-	return (*SystemEvent)(time.AfterFunc(d, f))
-}
-
-// SystemEvent implements Event using time.Timer.
-type SystemEvent time.Timer
-
-// Cancel implements Event.
-func (e *SystemEvent) Cancel() bool {
-	return (*time.Timer)(e).Stop()
+// AfterFunc runs f on a new goroutine after the duration has elapsed.
+func (System) AfterFunc(d time.Duration, f func()) Timer {
+	return time.AfterFunc(d, f)
 }
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
 }
diff --git a/common/mclock/simclock_test.go b/common/mclock/simclock_test.go
new file mode 100644
index 000000000..09e4391c1
--- /dev/null
+++ b/common/mclock/simclock_test.go
@@ -0,0 +1,115 @@
+// Copyright 2018 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package mclock
+
+import (
+	"testing"
+	"time"
+)
+
+var _ Clock = System{}
+var _ Clock = new(Simulated)
+
+func TestSimulatedAfter(t *testing.T) {
+	const timeout = 30 * time.Minute
+	const adv = time.Minute
+
+	var (
+		c   Simulated
+		end = c.Now().Add(timeout)
+		ch  = c.After(timeout)
+	)
+	for c.Now() < end.Add(-adv) {
+		c.Run(adv)
+		select {
+		case <-ch:
+			t.Fatal("Timer fired early")
+		default:
+		}
+	}
+
+	c.Run(adv)
+	select {
+	case stamp := <-ch:
+		want := time.Time{}.Add(timeout)
+		if !stamp.Equal(want) {
+			t.Errorf("Wrong time sent on timer channel: got %v, want %v", stamp, want)
+		}
+	default:
+		t.Fatal("Timer didn't fire")
+	}
+}
+
+func TestSimulatedAfterFunc(t *testing.T) {
+	var c Simulated
+
+	called1 := false
+	timer1 := c.AfterFunc(100*time.Millisecond, func() { called1 = true })
+	if c.ActiveTimers() != 1 {
+		t.Fatalf("%d active timers, want one", c.ActiveTimers())
+	}
+	if fired := timer1.Stop(); !fired {
+		t.Fatal("Stop returned false even though timer didn't fire")
+	}
+	if c.ActiveTimers() != 0 {
+		t.Fatalf("%d active timers, want zero", c.ActiveTimers())
+	}
+	if called1 {
+		t.Fatal("timer 1 called")
+	}
+	if fired := timer1.Stop(); fired {
+		t.Fatal("Stop returned true after timer was already stopped")
+	}
+
+	called2 := false
+	timer2 := c.AfterFunc(100*time.Millisecond, func() { called2 = true })
+	c.Run(50 * time.Millisecond)
+	if called2 {
+		t.Fatal("timer 2 called")
+	}
+	c.Run(51 * time.Millisecond)
+	if !called2 {
+		t.Fatal("timer 2 not called")
+	}
+	if fired := timer2.Stop(); fired {
+		t.Fatal("Stop returned true after timer has fired")
+	}
+}
+
+func TestSimulatedSleep(t *testing.T) {
+	var (
+		c       Simulated
+		timeout = 1 * time.Hour
+		done    = make(chan AbsTime)
+	)
+	go func() {
+		c.Sleep(timeout)
+		done <- c.Now()
+	}()
+
+	c.WaitForTimers(1)
+	c.Run(2 * timeout)
+	select {
+	case stamp := <-done:
+		want := AbsTime(2 * timeout)
+		if stamp != want {
+			t.Errorf("Wrong time after sleep: got %v, want %v", stamp, want)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("Sleep didn't return in time")
+	}
+}