add context cancel with reasone

Author: rekby

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added context cancel with specific error
+
 ## v3.26.10
 * Fixed syntax mistake in `trace.TablePooStateChangeInfo` to `trace.TablePoolStateChangeInfo` 
 

internal/conn/grpc_client_stream.go

@@ -2,11 +2,14 @@ package conn
 
 import (
 	"context"
+	"errors"
 	"time"
 
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 	"google.golang.org/grpc"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
+
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/wrap"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
@@ -40,7 +43,7 @@ func (s *grpcClientStream) CloseSend() (err error) {
 
 func (s *grpcClientStream) SendMsg(m interface{}) (err error) {
 	cancel := createPinger(s.c)
-	defer cancel()
+	defer cancel(xerrors.WithStackTrace(errors.New("send msg finished")))
 
 	err = s.ClientStream.SendMsg(m)
 
@@ -61,7 +64,7 @@ func (s *grpcClientStream) SendMsg(m interface{}) (err error) {
 
 func (s *grpcClientStream) RecvMsg(m interface{}) (err error) {
 	cancel := createPinger(s.c)
-	defer cancel()
+	defer cancel(xerrors.WithStackTrace(errors.New("receive msg finished")))
 
 	defer func() {
 		onDone := s.recv(xerrors.HideEOF(err))
@@ -102,9 +105,9 @@ func (s *grpcClientStream) RecvMsg(m interface{}) (err error) {
 	return nil
 }
 
-func createPinger(c *conn) context.CancelFunc {
+func createPinger(c *conn) xcontext.CancelErrFunc {
 	c.touchLastUsage()
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := xcontext.WithErrCancel(context.Background())
 	go func() {
 		ticker := time.NewTicker(time.Second)
 		ctxDone := ctx.Done()

internal/repeater/repeater.go

@@ -2,10 +2,15 @@ package repeater
 
 import (
 	"context"
+	"fmt"
 	"time"
 
 	"github.com/jonboulle/clockwork"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
+
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/backoff"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
@@ -27,7 +32,7 @@ type repeater struct {
 	// Task is a function that must be executed periodically.
 	task func(context.Context) error
 
-	cancel  context.CancelFunc
+	cancel  xcontext.CancelErrFunc
 	stopped chan struct{}
 
 	force chan struct{}
@@ -74,7 +79,7 @@ func New(
 	task func(ctx context.Context) (err error),
 	opts ...option,
 ) *repeater {
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := xcontext.WithErrCancel(context.Background())
 
 	r := &repeater{
 		interval: interval,
@@ -95,7 +100,7 @@ func New(
 }
 
 func (r *repeater) stop(onCancel func()) {
-	r.cancel()
+	r.cancel(xerrors.WithStackTrace(fmt.Errorf("ydb: repeater stopped")))
 	if onCancel != nil {
 		onCancel()
 	}

internal/table/session.go

@@ -2,6 +2,7 @@ package table
 
 import (
 	"context"
+	"fmt"
 	"net/url"
 	"strconv"
 	"sync"
@@ -11,6 +12,8 @@ import (
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/protobuf/proto"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
+
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer"
 
 	"github.com/ydb-platform/ydb-go-genproto/Ydb_Table_V1"
@@ -915,15 +918,15 @@ func (s *session) StreamReadTable(
 		opt((*options.ReadTableDesc)(&request))
 	}
 
-	ctx, cancel := context.WithCancel(ctx)
+	ctx, cancel := xcontext.WithErrCancel(ctx)
 
 	stream, err = s.tableService.StreamReadTable(
 		balancer.WithEndpoint(ctx, s),
 		&request,
 	)
 
 	if err != nil {
-		cancel()
+		cancel(xerrors.WithStackTrace(fmt.Errorf("ydb: stream read error: %w", err)))
 		return nil, xerrors.WithStackTrace(err)
 	}
 
@@ -950,7 +953,11 @@ func (s *session) StreamReadTable(
 			}
 		},
 		func(err error) error {
-			cancel()
+			if err == nil {
+				cancel(nil)
+			} else {
+				cancel(xerrors.WithStackTrace(fmt.Errorf("ydb: stream closed with: %w", err)))
+			}
 			onIntermediate(xerrors.HideEOF(err))(xerrors.HideEOF(err))
 			return err
 		},

internal/xcontext/cancel_with_error.go

@@ -2,6 +2,7 @@ package xcontext
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sync"
 	"time"
@@ -51,7 +52,6 @@ func (c *ctxError) errUnderLock() error {
 	}
 
 	return c.err
-
 }
 
 func (c *ctxError) Value(key interface{}) interface{} {
@@ -67,8 +67,29 @@ func (c *ctxError) cancel(err error) {
 	}
 
 	if c.errUnderLock() == nil {
+		err = cancelError{err: err}
 		c.err = err
 	}
 
 	c.ctxCancel()
 }
+
+type cancelError struct {
+	err error
+}
+
+func (e cancelError) Error() string {
+	return e.err.Error()
+}
+
+func (e cancelError) Is(target error) bool {
+	return errors.Is(e.err, target) || errors.Is(context.Canceled, target)
+}
+
+func (e cancelError) As(target interface{}) bool {
+	return errors.As(e.err, target) || errors.As(context.Canceled, target)
+}
+
+func (e cancelError) Unwrap() error {
+	return e.err
+}

internal/xcontext/cancel_with_error_test.go

@@ -13,7 +13,7 @@ func TestCancelWithError(t *testing.T) {
 	t.Run("SimpleCancel", func(t *testing.T) {
 		ctx, cancel := WithErrCancel(context.Background())
 		cancel(testError)
-		require.Equal(t, testError, ctx.Err())
+		require.ErrorIs(t, ctx.Err(), testError)
 	})
 
 	t.Run("CancelBeforeParent", func(t *testing.T) {
@@ -23,7 +23,8 @@ func TestCancelWithError(t *testing.T) {
 		cancel(testError)
 		parentCancel()
 
-		require.Equal(t, testError, ctx.Err())
+		require.ErrorIs(t, ctx.Err(), testError)
+		require.ErrorIs(t, ctx.Err(), context.Canceled)
 	})
 
 	t.Run("CancelAfterParent", func(t *testing.T) {
@@ -36,4 +37,10 @@ func TestCancelWithError(t *testing.T) {
 		require.Equal(t, context.Canceled, ctx.Err())
 	})
 
+	t.Run("CancelWithNil", func(t *testing.T) {
+		ctx, cancel := WithErrCancel(context.Background())
+		cancel(nil)
+		require.ErrorIs(t, ctx.Err(), errCancelWithNilError)
+		require.ErrorIs(t, ctx.Err(), context.Canceled)
+	})
 }