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..5cc6d80 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,32 @@ 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() +} + +func (c *Clock) ContextWithDeadline(ctx context.Context, t time.Time) (context.Context, context.CancelFunc) { + return c.ContextWithDeadlineCause(ctx, t, nil) +} + +func (c *Clock) ContextWithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) { + return c.ContextWithDeadlineCause(ctx, c.Now().Add(d), nil) +} + +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..2e27c95 --- /dev/null +++ b/fake/fake_clock_121.go @@ -0,0 +1,35 @@ +//go:build go1.20 + +package fake + +import ( + "context" + "time" +) + +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.CompareAndSwap(false, true) + cancelCause(cause) + return dctx, func() { + cancelCause(context.Canceled) + } + } + stop := f.AfterFunc(dur, func() { + if cctx.Err() == nil { + dctx.timedOut.CompareAndSwap(false, 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..e7074be --- /dev/null +++ b/fake/fake_clock_pre121.go @@ -0,0 +1,35 @@ +//go:build !go1.20 + +package fake + +import ( + "context" + "time" +) + +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.CompareAndSwap(false, true) + cancel() + return dctx, func() { + cancel() + } + } + stop := f.AfterFunc(dur, func() { + if cctx.Err() == nil { + dctx.timedOut.CompareAndSwap(false, 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..bf96d14 100644 --- a/offset/offset_clock.go +++ b/offset/offset_clock.go @@ -49,6 +49,14 @@ func (o *Clock) AfterFunc(d time.Duration, f func()) clocks.StopTimer { return o.inner.AfterFunc(d, f) } +func (o *Clock) ContextWithDeadline(ctx context.Context, t time.Time) (context.Context, context.CancelFunc) { + return o.inner.ContextWithDeadline(ctx, t.Add(o.offset)) +} + +func (o *Clock) ContextWithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) { + return o.inner.ContextWithTimeout(ctx, d+o.offset) +} + // 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..fd4f1b6 --- /dev/null +++ b/offset/offset_clock_121.go @@ -0,0 +1,16 @@ +//go:build go1.21 + +package offset + +import ( + "context" + "time" +) + +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) +} + +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..e3c7dd5 --- /dev/null +++ b/offset/offset_clock_pre121.go @@ -0,0 +1,16 @@ +//go:build !go1.21 + +package offset + +import ( + "context" + "time" +) + +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)) +} + +func (o *Clock) ContextWithTimeoutCause(ctx context.Context, d time.Duration, cause error) (context.Context, context.CancelFunc) { + return o.inner.ContextWithTimeout(ctx, d+o.offset) +}