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

* Fixed bug in `internal/xcontext.WithDone` (not listening chan done)

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed bug in `internal/xcontext.WithDone` (not listening chan done)
+
 ## v3.100.1
 * Refactored behaviour on `retry.Retryable` error for retry object (such as session, connection or transaction)
 

internal/xcontext/done.go

@@ -2,40 +2,25 @@ 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 {
+func WithDone(parent context.Context, done <-chan struct{}) (context.Context, context.CancelFunc) {
+	ctx, cancel := WithCancel(parent)
 	select {
 	case <-done:
-		return context.Canceled
+		cancel()
+
+		return ctx, cancel
 	default:
-		return nil
 	}
-}
 
-func (done doneCtx) Value(key any) any {
-	return nil
-}
+	go func() {
+		select {
+		case <-done:
+			cancel()
+		case <-parent.Done():
+		}
+	}()
 
-func WithDone(parent context.Context, done <-chan struct{}) (context.Context, context.CancelFunc) {
-	ctx, cancel := context.WithCancel(parent)
-	stop := context.AfterFunc(doneCtx(done), func() {
-		cancel()
-	})
-
-	return ctx, func() {
-		stop()
-		cancel()
-	}
+	return ctx, cancel
 }

internal/xcontext/done_test.go

@@ -16,7 +16,7 @@ func TestWithDone(t *testing.T) {
 		cancel()
 		require.Error(t, ctx1.Err())
 	})
-	t.Run("CloseDone", func(t *testing.T) {
+	t.Run("StopWithDone", func(t *testing.T) {
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
 		done := make(chan struct{})
@@ -26,4 +26,27 @@ func TestWithDone(t *testing.T) {
 		require.NoError(t, ctx.Err())
 		require.Error(t, ctx1.Err())
 	})
+	t.Run("CloseDone", func(t *testing.T) {
+		done := make(chan struct{})
+		ctx, cancel := WithDone(context.Background(), done)
+		require.NoError(t, ctx.Err())
+		close(done)
+		<-ctx.Done()
+		require.Error(t, ctx.Err())
+		cancel()
+	})
+	t.Run("ClosedDone", 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) {
+		var done chan struct{}
+		ctx, cancel := WithDone(context.Background(), done)
+		require.NoError(t, ctx.Err())
+		cancel()
+		require.Error(t, ctx.Err())
+	})
 }