go-ethereum/swarm/storage/feed/lookup/timesim_test.go

package lookup_test

// This file contains simple time simulation tools for testing
// and measuring time-aware algorithms

import (
	"sync"
	"time"
)

// Timer tracks information about a simulated timer
type Timer struct {
	deadline time.Time
	signal   chan time.Time
	id       int
}

// Stopwatch measures simulated execution time and manages simulated timers
type Stopwatch struct {
	t            time.Time
	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{
		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[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.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
			close(timer.signal)
			delete(s.timers, id)
		}
	}
}

// NewTimer returns a new timer that will trigger after `duration` elapses in the
// simulation
func (s *Stopwatch) NewTimer(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:       s.timerCounter,
	}

	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.NewTimer(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):
				s.Tick()
			case <-stopSignal:
				return
			}
		}
	}()
}

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