Merge pull request #1667 from ydb-platform/with-done

Reverted implementation of internal/xcontext.WithDone

Author: asmyasnikov

internal/xcontext/done.go

@@ -2,10 +2,35 @@ package xcontext
 
 import (
 	"context"
+	"time"
 )
 
+type doneCtx <-chan struct{}
+
+func (done doneCtx) Deadline() (deadline time.Time, ok bool) {
+	return
+}
+
+func (done doneCtx) Done() <-chan struct{} {
+	return done
+}
+
+func (done doneCtx) Err() error {
+	select {
+	case <-done:
+		return context.Canceled
+	default:
+		return nil
+	}
+}
+
+func (done doneCtx) Value(key any) any {
+	return nil
+}
+
 func WithDone(parent context.Context, done <-chan struct{}) (context.Context, context.CancelFunc) {
-	ctx, cancel := WithCancel(parent)
+	ctx, cancel := context.WithCancel(parent)
+
 	select {
 	case <-done:
 		cancel()
@@ -14,13 +39,12 @@ func WithDone(parent context.Context, done <-chan struct{}) (context.Context, co
 	default:
 	}
 
-	go func() {
-		select {
-		case <-done:
-			cancel()
-		case <-parent.Done():
-		}
-	}()
+	stop := context.AfterFunc(doneCtx(done), func() {
+		cancel()
+	})
 
-	return ctx, cancel
+	return ctx, func() {
+		stop()
+		cancel()
+	}
 }

internal/xcontext/done_test.go

@@ -8,15 +8,15 @@ import (
 )
 
 func TestWithDone(t *testing.T) {
-	t.Run("CancelParent", func(t *testing.T) {
+	t.Run("WithParentCancel", func(t *testing.T) {
 		ctx, cancel := context.WithCancel(context.Background())
 		done := make(chan struct{})
 		ctx1, _ := WithDone(ctx, done)
 		require.NoError(t, ctx1.Err())
 		cancel()
 		require.Error(t, ctx1.Err())
 	})
-	t.Run("StopWithDone", func(t *testing.T) {
+	t.Run("WithExplicitCancel", func(t *testing.T) {
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
 		done := make(chan struct{})
@@ -26,7 +26,7 @@ func TestWithDone(t *testing.T) {
 		require.NoError(t, ctx.Err())
 		require.Error(t, ctx1.Err())
 	})
-	t.Run("CloseDone", func(t *testing.T) {
+	t.Run("WithExplicitCloseDone", func(t *testing.T) {
 		done := make(chan struct{})
 		ctx, cancel := WithDone(context.Background(), done)
 		require.NoError(t, ctx.Err())
@@ -35,14 +35,14 @@ func TestWithDone(t *testing.T) {
 		require.Error(t, ctx.Err())
 		cancel()
 	})
-	t.Run("ClosedDone", func(t *testing.T) {
+	t.Run("WithClosedDone", func(t *testing.T) {
 		done := make(chan struct{})
 		close(done)
 		ctx, cancel := WithDone(context.Background(), done)
 		require.Error(t, ctx.Err())
 		cancel()
 	})
-	t.Run("NilDone", func(t *testing.T) {
+	t.Run("WithNilDone", func(t *testing.T) {
 		var done chan struct{}
 		ctx, cancel := WithDone(context.Background(), done)
 		require.NoError(t, ctx.Err())