diff --git a/internal/xcontext/done.go b/internal/xcontext/done.go index 762764d4e..7121d00a2 100644 --- a/internal/xcontext/done.go +++ b/internal/xcontext/done.go @@ -5,7 +5,32 @@ import ( "time" ) -type doneCtx <-chan struct{} +type ( + doneCtx <-chan struct{} + doneAlreadySignaledCtx struct { + context.Context //nolint:containedctx // thin wrapper delegating Deadline/Value + + err error + } +) + +var ( + noopCancel = func() {} + closedDoneChan = func() <-chan struct{} { + ch := make(chan struct{}) + close(ch) + + return ch + }() +) + +func (doneAlreadySignaledCtx) Done() <-chan struct{} { + return closedDoneChan +} + +func (ctx doneAlreadySignaledCtx) Err() error { + return ctx.err +} func (done doneCtx) Deadline() (deadline time.Time, ok bool) { return @@ -29,22 +54,31 @@ func (done doneCtx) Value(key any) any { } func WithDone(parent context.Context, done <-chan struct{}) (context.Context, context.CancelFunc) { - ctx, cancel := context.WithCancel(parent) + if parent.Err() != nil { + return parent, noopCancel + } select { case <-done: - cancel() + err := parent.Err() + if err == nil { + err = context.Canceled + } - return ctx, cancel + return doneAlreadySignaledCtx{ + Context: parent, + err: err, + }, noopCancel default: - } + ctx, cancel := context.WithCancel(parent) - stop := context.AfterFunc(doneCtx(done), func() { - cancel() - }) + stop := context.AfterFunc(doneCtx(done), func() { + cancel() + }) - return ctx, func() { - stop() - cancel() + return ctx, func() { + stop() + cancel() + } } } diff --git a/internal/xcontext/done_test.go b/internal/xcontext/done_test.go index 3fed30a2d..f0c15a48f 100644 --- a/internal/xcontext/done_test.go +++ b/internal/xcontext/done_test.go @@ -3,6 +3,7 @@ package xcontext import ( "context" "testing" + "time" "github.com/stretchr/testify/require" ) @@ -41,6 +42,21 @@ func TestWithDone(t *testing.T) { require.Error(t, ctx.Err()) cancel() }) + t.Run("WithClosedDoneErrShouldNotChangeAfterParentTimeout", func(t *testing.T) { + done := make(chan struct{}) + close(done) + + parent, cancelParent := context.WithTimeout(context.Background(), time.Millisecond) + t.Cleanup(cancelParent) + + ctx, cancel := WithDone(parent, done) + t.Cleanup(cancel) + require.ErrorIs(t, ctx.Err(), context.Canceled) + + <-parent.Done() + require.ErrorIs(t, parent.Err(), context.DeadlineExceeded) + require.ErrorIs(t, ctx.Err(), context.Canceled) + }) t.Run("WithNilDone", func(t *testing.T) { var done chan struct{} ctx, cancel := WithDone(context.Background(), done) @@ -49,3 +65,45 @@ func TestWithDone(t *testing.T) { require.Error(t, ctx.Err()) }) } + +// BenchmarkWithDone/AlreadyClosed-12 60969152 19.75 ns/op 16 B/op 1 allocs/op +// BenchmarkWithDone/Open_CancelImmediately-12 2763906 433.9 ns/op 312 B/op 7 allocs/op +// BenchmarkWithDone/Open_CloseDoneThenCancel-12 673682 1707 ns/op 536 B/op 9 allocs/op +func BenchmarkWithDone(b *testing.B) { + b.Run("AlreadyClosed", func(b *testing.B) { + done := make(chan struct{}) + close(done) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + ctx, cancel := WithDone(b.Context(), done) + cancel() + _ = ctx + } + }) + + b.Run("Open_CancelImmediately", func(b *testing.B) { + done := make(chan struct{}) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + ctx, cancel := WithDone(b.Context(), done) + cancel() + _ = ctx + } + }) + + b.Run("Open_CloseDoneThenCancel", func(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + b.StopTimer() + done := make(chan struct{}) + b.StartTimer() + ctx, cancel := WithDone(b.Context(), done) + close(done) + <-ctx.Done() + cancel() + } + }) +}