Skip to content

Commit

Permalink
Drop frozen dependency benbjohnson/clock
Browse files Browse the repository at this point in the history
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
  • Loading branch information
abhinav committed Sep 8, 2023
1 parent 2b35963 commit 0ec4e2b
Show file tree
Hide file tree
Showing 3 changed files with 147 additions and 13 deletions.
1 change: 0 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 0 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
@@ -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=
Expand Down
157 changes: 147 additions & 10 deletions internal/ztest/clock.go
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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:

Check warning on line 77 in internal/ztest/clock.go

View check run for this annotation

Codecov / codecov/patch

internal/ztest/clock.go#L77

Added line #L77 was not covered by tests
// 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")

Check warning on line 96 in internal/ztest/clock.go

View check run for this annotation

Codecov / codecov/patch

internal/ztest/clock.go#L96

Added line #L96 was not covered by tests
}

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
}

0 comments on commit 0ec4e2b

Please sign in to comment.