httpx/internal/clock/clock.go

package clock

import (
	"sync"
	"time"
)

// Clock abstracts time operations for deterministic testing.
type Clock interface {
	Now() time.Time
	Since(t time.Time) time.Duration
	NewTimer(d time.Duration) Timer
	After(d time.Duration) <-chan time.Time
}

// Timer abstracts time.Timer for testability.
type Timer interface {
	C() <-chan time.Time
	Stop() bool
	Reset(d time.Duration) bool
}

// System returns a Clock backed by the real system time.
func System() Clock { return systemClock{} }

type systemClock struct{}

func (systemClock) Now() time.Time                         { return time.Now() }
func (systemClock) Since(t time.Time) time.Duration        { return time.Since(t) }
func (systemClock) NewTimer(d time.Duration) Timer         { return &systemTimer{t: time.NewTimer(d)} }
func (systemClock) After(d time.Duration) <-chan time.Time { return time.After(d) }

type systemTimer struct{ t *time.Timer }

func (s *systemTimer) C() <-chan time.Time        { return s.t.C }
func (s *systemTimer) Stop() bool                 { return s.t.Stop() }
func (s *systemTimer) Reset(d time.Duration) bool { return s.t.Reset(d) }

// Mock returns a manually-controlled Clock for tests.
func Mock(now time.Time) *MockClock {
	return &MockClock{now: now}
}

// MockClock is a deterministic clock for testing.
type MockClock struct {
	mu     sync.Mutex
	now    time.Time
	timers []*mockTimer
}

func (m *MockClock) Now() time.Time {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.now
}

func (m *MockClock) Since(t time.Time) time.Duration {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.now.Sub(t)
}

func (m *MockClock) NewTimer(d time.Duration) Timer {
	m.mu.Lock()
	defer m.mu.Unlock()
	t := &mockTimer{
		clock:    m,
		ch:       make(chan time.Time, 1),
		deadline: m.now.Add(d),
		active:   true,
	}
	m.timers = append(m.timers, t)
	if d <= 0 {
		t.fire(m.now)
	}
	return t
}

func (m *MockClock) After(d time.Duration) <-chan time.Time {
	return m.NewTimer(d).C()
}

// Advance moves the clock forward by d and fires any expired timers.
func (m *MockClock) Advance(d time.Duration) {
	m.mu.Lock()
	m.now = m.now.Add(d)
	now := m.now
	m.mu.Unlock()

	m.fireExpired(now)
}

// Set sets the clock to an absolute time and fires any expired timers.
func (m *MockClock) Set(t time.Time) {
	m.mu.Lock()
	m.now = t
	now := m.now
	m.mu.Unlock()

	m.fireExpired(now)
}

// fireExpired fires all active timers whose deadline has passed, then
// removes inactive timers to prevent unbounded growth.
func (m *MockClock) fireExpired(now time.Time) {
	m.mu.Lock()
	timers := m.timers
	m.mu.Unlock()

	for _, t := range timers {
		t.mu.Lock()
		if t.active && !now.Before(t.deadline) {
			t.fire(now)
		}
		t.mu.Unlock()
	}

	// Compact: remove inactive timers. Use a new slice to avoid aliasing
	// the backing array (NewTimer may have appended between snapshots).
	m.mu.Lock()
	n := len(m.timers)
	active := make([]*mockTimer, 0, n)
	for _, t := range m.timers {
		t.mu.Lock()
		keep := t.active
		t.mu.Unlock()
		if keep {
			active = append(active, t)
		}
	}
	m.timers = active
	m.mu.Unlock()
}

type mockTimer struct {
	mu       sync.Mutex
	clock    *MockClock
	ch       chan time.Time
	deadline time.Time
	active   bool
}

func (t *mockTimer) C() <-chan time.Time { return t.ch }

func (t *mockTimer) Stop() bool {
	t.mu.Lock()
	defer t.mu.Unlock()
	was := t.active
	t.active = false
	return was
}

func (t *mockTimer) Reset(d time.Duration) bool {
	// Acquire clock lock first to match the lock ordering in fireExpired
	// (clock.mu → t.mu), preventing deadlock.
	t.clock.mu.Lock()
	deadline := t.clock.now.Add(d)
	t.clock.mu.Unlock()

	t.mu.Lock()
	defer t.mu.Unlock()
	was := t.active
	t.active = true
	t.deadline = deadline
	return was
}

// fire sends the time on the channel. Caller must hold t.mu.
func (t *mockTimer) fire(now time.Time) {
	t.active = false
	select {
	case t.ch <- now:
	default:
	}
}