swarm/storage/feed/lookup: Add comments to tests

swarm/storage/feed/lookup/timesim_test.go

@@ -1,44 +1,54 @@
 package lookup_test
 
+// This file contains simple time simulation tools for testing
+// and measuring time-aware algorithms
+
 import (
 	"sync"
-	"sync/atomic"
 	"time"
 )
 
+// Timer tracks information about a simulated timer
 type Timer struct {
 	deadline time.Time
 	signal   chan time.Time
-	id       int32
+	id       int
 }
 
+// Stopwatch measures simulated execution time and manages simulated timers
 type Stopwatch struct {
 	t            time.Time
-	r            time.Duration
-	timers       map[int32]*Timer
-	timerCounter int32
+	resolution   time.Duration
+	timers       map[int]*Timer
+	timerCounter int
 	stopSignal   chan struct{}
 	lock         sync.RWMutex
 }
 
+// NewStopwatch returns a simulated clock that ticks on `resolution` intervals
 func NewStopwatch(resolution time.Duration) *Stopwatch {
 	s := &Stopwatch{
-		r: resolution,
+		resolution: resolution,
 	}
 	s.Reset()
 	return s
 }
 
+// Reset clears all timers and sents the stopwatch to zero
 func (s *Stopwatch) Reset() {
 	s.t = time.Time{}
-	s.timers = make(map[int32]*Timer)
+	s.timers = make(map[int]*Timer)
 	s.Stop()
 }
 
+// Tick advances simulated time by the stopwatch's resolution and triggers
+// all due timers
 func (s *Stopwatch) Tick() {
-	s.t = s.t.Add(s.r)
+	s.t = s.t.Add(s.resolution)
+
 	s.lock.Lock()
 	defer s.lock.Unlock()
+
 	for id, timer := range s.timers {
 		if s.t.After(timer.deadline) || s.t.Equal(timer.deadline) {
 			timer.signal <- s.t
@@ -48,33 +58,45 @@ func (s *Stopwatch) Tick() {
 	}
 }
 
+// GetTimer returns a new timer that will trigger after `duration` elapses in the
+// simulation
 func (s *Stopwatch) GetTimer(duration time.Duration) <-chan time.Time {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	s.timerCounter++
 	timer := &Timer{
 		deadline: s.t.Add(duration),
 		signal:   make(chan time.Time, 1),
-		id:       atomic.AddInt32(&s.timerCounter, 1),
+		id:       s.timerCounter,
 	}
-	s.lock.Lock()
-	defer s.lock.Unlock()
 
 	s.timers[timer.id] = timer
 	return timer.signal
 }
 
+// TimeAfter returns a simulated timer factory that can replace `time.After`
 func (s *Stopwatch) TimeAfter() func(d time.Duration) <-chan time.Time {
 	return func(d time.Duration) <-chan time.Time {
 		return s.GetTimer(d)
 	}
 }
 
+// Elapsed returns the time that has passed in the simulation
 func (s *Stopwatch) Elapsed() time.Duration {
 	return s.t.Sub(time.Time{})
 }
 
+// Run starts the time simulation
 func (s *Stopwatch) Run() {
 	go func() {
 		stopSignal := make(chan struct{})
+		s.lock.Lock()
+		if s.stopSignal != nil {
+			close(s.stopSignal)
+		}
 		s.stopSignal = stopSignal
+		s.lock.Unlock()
 		for {
 			select {
 			case <-time.After(1 * time.Millisecond):
@@ -86,9 +108,21 @@ func (s *Stopwatch) Run() {
 	}()
 }
 
+// Stop stops the time simulation
 func (s *Stopwatch) Stop() {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
 	if s.stopSignal != nil {
 		close(s.stopSignal)
 		s.stopSignal = nil
 	}
 }
+
+func (s *Stopwatch) Measure(measuredFunc func()) time.Duration {
+	s.Reset()
+	s.Run()
+	defer s.Stop()
+	measuredFunc()
+	return s.Elapsed()
+}