nmon/internal/clock/clock.go

// Package clock is the only source of time for nmon logic. Production code
// uses Real(); tests use NewFake so adversarial fixtures never wait on real time.
package clock

import (
	"context"
	"sort"
	"sync"
	"time"
)

// Clock abstracts time so schedulers, deadlines and rollups are testable.
type Clock interface {
	Now() time.Time
	After(d time.Duration) <-chan time.Time
	NewTicker(d time.Duration) Ticker
	// Sleep waits for d or until ctx is done, returning ctx.Err() in the latter case.
	Sleep(ctx context.Context, d time.Duration) error
}

// Ticker mirrors time.Ticker.
type Ticker interface {
	C() <-chan time.Time
	Stop()
}

// Real returns the wall clock.
func Real() Clock { return realClock{} }

type realClock struct{}

func (realClock) Now() time.Time                         { return time.Now() }
func (realClock) After(d time.Duration) <-chan time.Time { return time.After(d) }
func (realClock) NewTicker(d time.Duration) Ticker       { return realTicker{time.NewTicker(d)} }
func (realClock) Sleep(ctx context.Context, d time.Duration) error {
	t := time.NewTimer(d)
	defer t.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-t.C:
		return nil
	}
}

type realTicker struct{ t *time.Ticker }

func (r realTicker) C() <-chan time.Time { return r.t.C }
func (r realTicker) Stop()               { r.t.Stop() }

// Fake is a manually advanced clock. Timers fire when Advance or Set moves
// time past their deadline.
type Fake struct {
	mu      sync.Mutex
	now     time.Time
	waiters []*waiter
}

type waiter struct {
	at     time.Time
	ch     chan time.Time
	period time.Duration // >0 for tickers
	done   bool
}

// NewFake returns a Fake clock starting at start.
func NewFake(start time.Time) *Fake { return &Fake{now: start} }

// Now returns the fake time.
func (f *Fake) Now() time.Time {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.now
}

// After returns a channel that fires once Advance moves the fake time past d.
func (f *Fake) After(d time.Duration) <-chan time.Time {
	f.mu.Lock()
	defer f.mu.Unlock()
	w := &waiter{at: f.now.Add(d), ch: make(chan time.Time, 1)}
	if d <= 0 {
		w.ch <- f.now
		return w.ch
	}
	f.waiters = append(f.waiters, w)
	return w.ch
}

// NewTicker returns a ticker that fires each time Advance crosses a multiple of d.
func (f *Fake) NewTicker(d time.Duration) Ticker {
	if d <= 0 {
		panic("clock: non-positive ticker interval")
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	w := &waiter{at: f.now.Add(d), ch: make(chan time.Time, 1), period: d}
	f.waiters = append(f.waiters, w)
	return &fakeTicker{f: f, w: w}
}

// Sleep blocks until Advance moves the fake time past d, or ctx is done.
func (f *Fake) Sleep(ctx context.Context, d time.Duration) error {
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-f.After(d):
		return nil
	}
}

// Advance moves the clock forward by d, firing every timer due on the way in order.
func (f *Fake) Advance(d time.Duration) { f.Set(f.Now().Add(d)) }

// Set moves the clock to t (never backwards), firing due timers in order.
func (f *Fake) Set(t time.Time) {
	for {
		f.mu.Lock()
		if t.Before(f.now) {
			f.mu.Unlock()
			return
		}
		var live []*waiter
		for _, w := range f.waiters {
			if !w.done {
				live = append(live, w)
			}
		}
		f.waiters = live
		sort.Slice(live, func(i, j int) bool { return live[i].at.Before(live[j].at) })
		if len(live) == 0 || live[0].at.After(t) {
			f.now = t
			f.mu.Unlock()
			return
		}
		w := live[0]
		f.now = w.at
		select {
		case w.ch <- w.at:
		default: // ticker semantics: drop if the receiver is behind
		}
		if w.period > 0 {
			w.at = w.at.Add(w.period)
		} else {
			w.done = true
		}
		f.mu.Unlock()
	}
}

// Waiters reports how many timers are pending; tests use it to wait until a
// goroutine has blocked on the clock before advancing.
func (f *Fake) Waiters() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	n := 0
	for _, w := range f.waiters {
		if !w.done {
			n++
		}
	}
	return n
}

type fakeTicker struct {
	f *Fake
	w *waiter
}

func (t *fakeTicker) C() <-chan time.Time { return t.w.ch }
func (t *fakeTicker) Stop() {
	t.f.mu.Lock()
	t.w.done = true
	t.f.mu.Unlock()
}