diff --git a/common/mclock/mclock.go b/common/mclock/mclock.go index 76d1e5cdb7..c07d55de93 100644 --- a/common/mclock/mclock.go +++ b/common/mclock/mclock.go @@ -33,15 +33,24 @@ func Now() AbsTime { // Clock interface makes it possible to replace the monotonic system clock with // a simulated clock +// +// Note: event loops capable of running with a simulated clock should listen to PingChannel. +// MonotonicClock also implements this function to ensure interface compatibility. type Clock interface { Now() AbsTime Sleep(time.Duration) After(time.Duration) <-chan time.Time + PingChannel() chan struct{} } // MonotonicClock implements Clock using the system clock type MonotonicClock struct{} +// PingChannel implements Clock by returning a dummy nil channel +func (MonotonicClock) PingChannel() chan struct{} { + return nil +} + // Now implements Clock func (MonotonicClock) Now() AbsTime { return AbsTime(monotime.Now()) diff --git a/common/mclock/simclock.go b/common/mclock/simclock.go index 11cff0dc4a..2030fec906 100644 --- a/common/mclock/simclock.go +++ b/common/mclock/simclock.go @@ -18,7 +18,6 @@ package mclock import ( - "runtime" "sync" "time" ) @@ -39,18 +38,23 @@ type SimulatedClock struct { now AbsTime scheduled []event stop bool + pingCh chan struct{} lock sync.RWMutex } // NewSimulatedClock creates a new simulated clock -func NewSimulatedClock(maxStep time.Duration, goSchedCount int) *SimulatedClock { - s := &SimulatedClock{scheduled: make([]event, 0, 100)} +func NewSimulatedClock(maxStep time.Duration, pingCount int) *SimulatedClock { + s := &SimulatedClock{scheduled: make([]event, 0, 100), pingCh: make(chan struct{})} go func() { lastScheduled := 0 for { - for i := 0; i < goSchedCount; i++ { - runtime.Gosched() + timeout := time.After(maxStep / 100) + for i := 0; i < pingCount; i++ { + select { + case s.pingCh <- struct{}{}: + case <-timeout: + } } s.lock.Lock() if s.stop { @@ -76,6 +80,11 @@ func NewSimulatedClock(maxStep time.Duration, goSchedCount int) *SimulatedClock return s } +// PingChannel returns a channel that event loops should listen to +func (s *SimulatedClock) PingChannel() chan struct{} { + return s.pingCh +} + // Stop stops the clock (Sleeps and Afters will never return after this) func (s *SimulatedClock) Stop() { s.lock.Lock() @@ -97,7 +106,13 @@ func (s *SimulatedClock) Sleep(d time.Duration) { s.insert(d, func() { close(done) }) - <-done + for { + select { + case <-done: + return + case <-s.pingCh: + } + } } // After implements Clock diff --git a/les/freeclient_test.go b/les/freeclient_test.go index f863cc0f90..1ca9750d2d 100644 --- a/les/freeclient_test.go +++ b/les/freeclient_test.go @@ -86,6 +86,7 @@ func testFreeClientPool(t *testing.T, connLimit, clientCount int) { loop: for { select { + case <-clock.PingChannel(): case <-tickCh: i := rand.Intn(clientCount) if connected[i] {