diff --git a/go.mod b/go.mod index 455dae496..9a091d941 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,6 @@ module go.uber.org/zap go 1.19 require ( - github.com/benbjohnson/clock v1.3.0 github.com/stretchr/testify v1.8.1 go.uber.org/goleak v1.2.0 go.uber.org/multierr v1.10.0 diff --git a/go.sum b/go.sum index ffa795531..6f3b5b06c 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,3 @@ -github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A= -github.com/benbjohnson/clock v1.3.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= diff --git a/internal/ztest/clock.go b/internal/ztest/clock.go index fe8026d94..cfdd83d69 100644 --- a/internal/ztest/clock.go +++ b/internal/ztest/clock.go @@ -1,4 +1,4 @@ -// Copyright (c) 2021 Uber Technologies, Inc. +// Copyright (c) 2023 Uber Technologies, Inc. // // Permission is hereby granted, free of charge, to any person obtaining a copy // of this software and associated documentation files (the "Software"), to deal @@ -21,30 +21,167 @@ package ztest import ( + "container/heap" + "sync" "time" - - "github.com/benbjohnson/clock" ) -// MockClock provides control over the time. -type MockClock struct{ m *clock.Mock } +// MockClock is a fake source of time. +// It implements standard time operations, +// but allows the user to control the passage of time. +// +// Use the [Add] method to progress time. +type MockClock struct { + mu sync.RWMutex + now time.Time -// NewMockClock builds a new mock clock that provides control of time. + // The MockClock works by maintaining a list of waiters. + // Each waiter knows the time at which it should be resolved. + // When the clock advances, all waiters that are in range are resolved + // in chronological order. + waiters waiters +} + +// NewMockClock builds a new mock clock +// using the current actual time as the initial time. func NewMockClock() *MockClock { - return &MockClock{clock.NewMock()} + return &MockClock{ + now: time.Now(), + } } // Now reports the current time. func (c *MockClock) Now() time.Time { - return c.m.Now() + c.mu.RLock() + defer c.mu.RUnlock() + return c.now } // NewTicker returns a time.Ticker that ticks at the specified frequency. +// +// As with [time.NewTicker], +// the ticker will drop ticks if the receiver is slow, +// and the channel is never closed. func (c *MockClock) NewTicker(d time.Duration) *time.Ticker { - return &time.Ticker{C: c.m.Ticker(d).C} + ch := make(chan time.Time, 1) + + var tick func(time.Time) + tick = func(now time.Time) { + next := now.Add(d) + c.runAt(next, func() { + defer tick(next) + + select { + case ch <- next: + // ok + default: + // The receiver is slow. + // Drop the tick and continue. + } + }) + } + tick(c.Now()) + + return &time.Ticker{C: ch} } // Add progresses time by the given duration. +// +// Other operations waiting for the time to advance +// will be resolved if they are within range. +// +// Panics if the duration is negative. func (c *MockClock) Add(d time.Duration) { - c.m.Add(d) + if d < 0 { + panic("cannot add negative duration") + } + + c.mu.Lock() + defer c.mu.Unlock() + + newTime := c.now.Add(d) + // newTime won't be recorded until the end of this method. + // This ensures that any waiters that are resolved + // are resolved at the time they were expecting. + + for w, ok := c.waiters.PopLTE(newTime); ok; w, ok = c.waiters.PopLTE(newTime) { + // The waiter is within range. + // Travel to the time of the waiter and resolve it. + c.now = w.until + + // The waiter may schedule more work + // so we must release the lock. + c.mu.Unlock() + w.fn() + // Sleeping here is necessary to let the side effects of waiters + // take effect before we continue. + time.Sleep(1 * time.Millisecond) + c.mu.Lock() + } + + c.now = newTime +} + +// runAt schedules the given function to be run at the given time. +// The function runs without a lock held, so it may schedule more work. +func (c *MockClock) runAt(t time.Time, fn func()) { + c.mu.Lock() + defer c.mu.Unlock() + c.waiters.Push(waiter{until: t, fn: fn}) +} + +type waiter struct { + until time.Time + fn func() +} + +// waiters is a thread-safe collection of waiters +// with the next waiter to be resolved at the front. +// +// Use the methods on this type to manipulate the collection. +// Do not modify the slice directly. +type waiters struct{ heap waiterHeap } + +// Push adds a new waiter to the collection. +func (w *waiters) Push(v waiter) { + heap.Push(&w.heap, v) +} + +// PopLTE removes and returns the next waiter to be resolved +// if it is scheduled to be resolved at or before the given time. +// +// Returns false if there are no waiters in range. +func (w *waiters) PopLTE(t time.Time) (_ waiter, ok bool) { + if len(w.heap) == 0 || w.heap[0].until.After(t) { + return waiter{}, false + } + + return heap.Pop(&w.heap).(waiter), true +} + +// waiterHeap implements a min-heap of waiters based on their 'until' time. +// +// This is separate from the waiters type so that we can implement heap.Interface +// while still exposing a type-safe API on waiters. +type waiterHeap []waiter + +var _ heap.Interface = (*waiterHeap)(nil) + +func (h waiterHeap) Len() int { return len(h) } +func (h waiterHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } + +func (h waiterHeap) Less(i, j int) bool { + return h[i].until.Before(h[j].until) +} + +func (h *waiterHeap) Push(x interface{}) { + *h = append(*h, x.(waiter)) +} + +func (h *waiterHeap) Pop() interface{} { + old := *h + n := len(old) + x := old[n-1] + *h = old[:n-1] + return x }