From 0f50b7f038c661daf58d8262a81fb0d9fdb36ade Mon Sep 17 00:00:00 2001 From: Justin Ruggles Date: Thu, 26 Oct 2023 09:33:37 -0400 Subject: [PATCH] Add context methods to Clock interface This allows for using timeout/deadline functionality built in to context.Context with a custom clock implementation. Module Go version bumped to 1.19 due to use of atomic.Bool --- clock.go | 23 ++++++++++++++++++++++ clock_121.go | 16 +++++++++++++++ clock_pre121.go | 16 +++++++++++++++ fake/fake_clock.go | 37 +++++++++++++++++++++++++++++++++++ fake/fake_clock_121.go | 36 ++++++++++++++++++++++++++++++++++ fake/fake_clock_pre121.go | 36 ++++++++++++++++++++++++++++++++++ go.mod | 2 +- offset/offset_clock.go | 13 ++++++++++++ offset/offset_clock_121.go | 22 +++++++++++++++++++++ offset/offset_clock_pre121.go | 22 +++++++++++++++++++++ 10 files changed, 222 insertions(+), 1 deletion(-) create mode 100644 clock_121.go create mode 100644 clock_pre121.go create mode 100644 fake/fake_clock_121.go create mode 100644 fake/fake_clock_pre121.go create mode 100644 offset/offset_clock_121.go create mode 100644 offset/offset_clock_pre121.go diff --git a/clock.go b/clock.go index 2ff2c2a..c5b5475 100644 --- a/clock.go +++ b/clock.go @@ -54,6 +54,14 @@ func (c defaultClock) AfterFunc(d time.Duration, f func()) StopTimer { return time.AfterFunc(d, f) } +func (c defaultClock) ContextWithDeadline(ctx context.Context, t time.Time) (context.Context, context.CancelFunc) { + return context.WithDeadline(ctx, t) +} + +func (c defaultClock) ContextWithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) { + return context.WithTimeout(ctx, d) +} + // DefaultClock returns a clock that minimally wraps the `time` package func DefaultClock() Clock { return defaultClock{} @@ -80,4 +88,19 @@ type Clock interface { // The callback function f will be executed after the interval d has // elapsed, unless the returned timer's Stop() method is called first. AfterFunc(d time.Duration, f func()) StopTimer + + // ContextWithDeadline behaves like context.WithDeadline, but it uses the + // clock to determine the when the deadline has expired. + ContextWithDeadline(ctx context.Context, t time.Time) (context.Context, context.CancelFunc) + // ContextWithDeadlineCause behaves like context.WithDeadlineCause, but it + // uses the clock to determine the when the deadline has expired. Cause is + // ignored in Go 1.20 and earlier. + ContextWithDeadlineCause(ctx context.Context, t time.Time, cause error) (context.Context, context.CancelFunc) + // ContextWithTimeout behaves like context.WithTimeout, but it uses the + // clock to determine the when the timeout has elapsed. + ContextWithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) + // ContextWithTimeoutCause behaves like context.WithTimeoutCause, but it + // uses the clock to determine the when the timeout has elapsed. Cause is + // ignored in Go 1.20 and earlier. + ContextWithTimeoutCause(ctx context.Context, d time.Duration, cause error) (context.Context, context.CancelFunc) } diff --git a/clock_121.go b/clock_121.go new file mode 100644 index 0000000..d1939d3 --- /dev/null +++ b/clock_121.go @@ -0,0 +1,16 @@ +//go:build go1.21 + +package clocks + +import ( + "context" + "time" +) + +func (c defaultClock) ContextWithDeadlineCause(ctx context.Context, t time.Time, cause error) (context.Context, context.CancelFunc) { + return context.WithDeadlineCause(ctx, t, cause) +} + +func (c defaultClock) ContextWithTimeoutCause(ctx context.Context, d time.Duration, cause error) (context.Context, context.CancelFunc) { + return context.WithTimeoutCause(ctx, d, cause) +} diff --git a/clock_pre121.go b/clock_pre121.go new file mode 100644 index 0000000..cc91707 --- /dev/null +++ b/clock_pre121.go @@ -0,0 +1,16 @@ +//go:build !go1.21 + +package clocks + +import ( + "context" + "time" +) + +func (c defaultClock) ContextWithDeadlineCause(ctx context.Context, t time.Time, cause error) (context.Context, context.CancelFunc) { + return context.WithDeadline(ctx, t) +} + +func (c defaultClock) ContextWithTimeoutCause(ctx context.Context, d time.Duration, cause error) (context.Context, context.CancelFunc) { + return context.WithTimeout(ctx, d) +} diff --git a/fake/fake_clock.go b/fake/fake_clock.go index 2b25c6d..093d666 100644 --- a/fake/fake_clock.go +++ b/fake/fake_clock.go @@ -3,6 +3,7 @@ package fake import ( "context" "sync" + "sync/atomic" "time" clocks "github.com/vimeo/go-clocks" @@ -410,3 +411,39 @@ func (f *Clock) AwaitTimerAborts(n int) { func (f *Clock) WaitAfterFuncs() { f.cbsWG.Wait() } + +type deadlineContext struct { + context.Context + timedOut atomic.Bool + deadline time.Time +} + +func (d *deadlineContext) Deadline() (time.Time, bool) { + return d.deadline, true +} + +func (d *deadlineContext) Err() error { + if d.timedOut.Load() { + return context.DeadlineExceeded + } + return d.Context.Err() +} + +// ContextWithDeadline behaves like context.WithDeadline, but it uses the +// clock to determine the when the deadline has expired. +func (c *Clock) ContextWithDeadline(ctx context.Context, t time.Time) (context.Context, context.CancelFunc) { + return c.ContextWithDeadlineCause(ctx, t, nil) +} + +// ContextWithTimeout behaves like context.WithTimeout, but it uses the +// clock to determine the when the timeout has elapsed. +func (c *Clock) ContextWithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) { + return c.ContextWithDeadlineCause(ctx, c.Now().Add(d), nil) +} + +// ContextWithTimeoutCause behaves like context.WithTimeoutCause, but it +// uses the clock to determine the when the timeout has elapsed. Cause is +// ignored in Go 1.20 and earlier. +func (c *Clock) ContextWithTimeoutCause(ctx context.Context, d time.Duration, cause error) (context.Context, context.CancelFunc) { + return c.ContextWithDeadlineCause(ctx, c.Now().Add(d), cause) +} diff --git a/fake/fake_clock_121.go b/fake/fake_clock_121.go new file mode 100644 index 0000000..4e68708 --- /dev/null +++ b/fake/fake_clock_121.go @@ -0,0 +1,36 @@ +//go:build go1.20 + +package fake + +import ( + "context" + "time" +) + +// ContextWithDeadlineCause behaves like context.WithDeadlineCause, but it +// uses the clock to determine the when the deadline has expired. Cause is +// ignored in Go 1.20 and earlier. +func (f *Clock) ContextWithDeadlineCause(ctx context.Context, t time.Time, cause error) (context.Context, context.CancelFunc) { + cctx, cancelCause := context.WithCancelCause(ctx) + dctx := &deadlineContext{ + Context: cctx, + deadline: t, + } + dur := f.Until(t) + if dur <= 0 { + dctx.timedOut.Store(true) + cancelCause(cause) + return dctx, func() {} + } + stop := f.AfterFunc(dur, func() { + if cctx.Err() == nil { + dctx.timedOut.Store(true) + } + cancelCause(cause) + }) + cancel := func() { + cancelCause(context.Canceled) + stop.Stop() + } + return dctx, cancel +} diff --git a/fake/fake_clock_pre121.go b/fake/fake_clock_pre121.go new file mode 100644 index 0000000..d5ad96e --- /dev/null +++ b/fake/fake_clock_pre121.go @@ -0,0 +1,36 @@ +//go:build !go1.20 + +package fake + +import ( + "context" + "time" +) + +// ContextWithDeadlineCause behaves like context.WithDeadlineCause, but it +// uses the clock to determine the when the deadline has expired. Cause is +// ignored in Go 1.20 and earlier. +func (f *Clock) ContextWithDeadlineCause(ctx context.Context, t time.Time, cause error) (context.Context, context.CancelFunc) { + cctx, cancel := context.WithCancel(ctx) + dctx := &deadlineContext{ + Context: cctx, + deadline: t, + } + dur := f.Until(t) + if dur <= 0 { + dctx.timedOut.Store(true) + cancel() + return dctx, func() {} + } + stop := f.AfterFunc(dur, func() { + if cctx.Err() == nil { + dctx.timedOut.Store(true) + } + cancel() + }) + cancelStop := func() { + cancel() + stop.Stop() + } + return dctx, cancelStop +} diff --git a/go.mod b/go.mod index 8790e9a..2bdc894 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,3 @@ module github.com/vimeo/go-clocks -go 1.14 +go 1.19 diff --git a/offset/offset_clock.go b/offset/offset_clock.go index a83375a..9a4ffea 100644 --- a/offset/offset_clock.go +++ b/offset/offset_clock.go @@ -49,6 +49,19 @@ func (o *Clock) AfterFunc(d time.Duration, f func()) clocks.StopTimer { return o.inner.AfterFunc(d, f) } +// ContextWithDeadline behaves like context.WithDeadline, but it uses the +// clock to determine the when the deadline has expired. +func (o *Clock) ContextWithDeadline(ctx context.Context, t time.Time) (context.Context, context.CancelFunc) { + return o.inner.ContextWithDeadline(ctx, t.Add(o.offset)) +} + +// ContextWithTimeout behaves like context.WithTimeout, but it uses the +// clock to determine the when the timeout has elapsed. +func (o *Clock) ContextWithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) { + // timeout is relative, so it doesn't need any adjustment + return o.inner.ContextWithTimeout(ctx, d) +} + // NewOffsetClock creates an OffsetClock. offset is added to all absolute times. func NewOffsetClock(inner clocks.Clock, offset time.Duration) *Clock { return &Clock{ diff --git a/offset/offset_clock_121.go b/offset/offset_clock_121.go new file mode 100644 index 0000000..dbeab92 --- /dev/null +++ b/offset/offset_clock_121.go @@ -0,0 +1,22 @@ +//go:build go1.21 + +package offset + +import ( + "context" + "time" +) + +// ContextWithDeadlineCause behaves like context.WithDeadlineCause, but it +// uses the clock to determine the when the deadline has expired. Cause is +// ignored in Go 1.20 and earlier. +func (o *Clock) ContextWithDeadlineCause(ctx context.Context, t time.Time, cause error) (context.Context, context.CancelFunc) { + return o.inner.ContextWithDeadlineCause(ctx, t.Add(o.offset), cause) +} + +// ContextWithTimeoutCause behaves like context.WithTimeoutCause, but it +// uses the clock to determine the when the timeout has elapsed. Cause is +// ignored in Go 1.20 and earlier. +func (o *Clock) ContextWithTimeoutCause(ctx context.Context, d time.Duration, cause error) (context.Context, context.CancelFunc) { + return o.inner.ContextWithTimeoutCause(ctx, d+o.offset, cause) +} diff --git a/offset/offset_clock_pre121.go b/offset/offset_clock_pre121.go new file mode 100644 index 0000000..df0f120 --- /dev/null +++ b/offset/offset_clock_pre121.go @@ -0,0 +1,22 @@ +//go:build !go1.21 + +package offset + +import ( + "context" + "time" +) + +// ContextWithDeadlineCause behaves like context.WithDeadlineCause, but it +// uses the clock to determine the when the deadline has expired. Cause is +// ignored in Go 1.20 and earlier. +func (o *Clock) ContextWithDeadlineCause(ctx context.Context, t time.Time, cause error) (context.Context, context.CancelFunc) { + return o.inner.ContextWithDeadline(ctx, t.Add(o.offset)) +} + +// ContextWithTimeoutCause behaves like context.WithTimeoutCause, but it +// uses the clock to determine the when the timeout has elapsed. Cause is +// ignored in Go 1.20 and earlier. +func (o *Clock) ContextWithTimeoutCause(ctx context.Context, d time.Duration, cause error) (context.Context, context.CancelFunc) { + return o.inner.ContextWithTimeout(ctx, d+o.offset) +}