From 0ec4e2b5aad6d850af0fe390c01e3e18c4d2df89 Mon Sep 17 00:00:00 2001 From: Abhinav Gupta Date: Mon, 21 Aug 2023 19:56:31 -0700 Subject: [PATCH] Drop frozen dependency benbjohnson/clock This drops the dependency on benbjohnson/clock in favor of a simple hand-rolled mock clock implementation. The core of the functionality of MockClock is provided by the following: - waiters: a min-heap of functions waiting to be executed - runAt: schedules a function to be executed when time progresses - Add: moves time forward, running all functions in range The rest of the time functionality can be built upon these pieces. Note that nothing happens until Add is called. There are no goroutines running additional work in the background. Resolves #1331 --- go.mod | 1 - go.sum | 2 - internal/ztest/clock.go | 157 +++++++++++++++++++++++++++++++++++++--- 3 files changed, 147 insertions(+), 13 deletions(-) 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 }