common/mclock: use ping channel instead of GoSched

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())

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 {
+func NewSimulatedClock(maxStep time.Duration, pingCount int) *SimulatedClock {
-	s := &SimulatedClock{scheduled: make([]event, 0, 100)}
+	s := &SimulatedClock{scheduled: make([]event, 0, 100), pingCh: make(chan struct{})}
 
 	go func() {
 		lastScheduled := 0
 		for {
-			for i := 0; i < goSchedCount; i++ {
+			timeout := time.After(maxStep / 100)
-				runtime.Gosched()
+			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

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] {