Merge branch 'master' into xsql-conn-close

Author: asmyasnikov

CHANGELOG.md

@@ -1,6 +1,8 @@
 * Fixed closing of `database/sql` connection (aka `YDB` session)
 * Made `session.Close()` as `nop` for idled session
 * Implemented goroutine for closing idle connection in `database/sql` driver
+* Separated errors of commit from other reader and to expired session
+* Fixed wrapping error in `internal/balancer/Balancer.wrapCall()`
 
 ## v3.42.4
 * Added `ydb.WithDisableServerBalancer()` database/sql connector option

internal/balancer/balancer.go

@@ -279,7 +279,10 @@ func (b *Balancer) wrapCall(ctx context.Context, f func(ctx context.Context, cc
 	}
 
 	if err = f(ctx, cc); err != nil {
-		return xerrors.WithStackTrace(err)
+		if conn.UseWrapping(ctx) {
+			return xerrors.WithStackTrace(err)
+		}
+		return err
 	}
 
 	return nil

internal/conn/conn.go

@@ -297,10 +297,10 @@ func (c *conn) Invoke(
 	opts ...grpc.CallOption,
 ) (err error) {
 	var (
-		opID     string
-		issues   []trace.Issue
-		wrapping = needWrapping(ctx)
-		onDone   = trace.DriverOnConnInvoke(
+		opID        string
+		issues      []trace.Issue
+		useWrapping = UseWrapping(ctx)
+		onDone      = trace.DriverOnConnInvoke(
 			c.config.Trace(),
 			&ctx,
 			c.endpoint,
@@ -330,7 +330,7 @@ func (c *conn) Invoke(
 
 	err = cc.Invoke(ctx, method, req, res, append(opts, grpc.Trailer(&md))...)
 	if err != nil {
-		if wrapping {
+		if useWrapping {
 			err = xerrors.FromGRPCError(err,
 				xerrors.WithAddress(c.Address()),
 			)
@@ -353,7 +353,7 @@ func (c *conn) Invoke(
 		for _, issue := range o.GetOperation().GetIssues() {
 			issues = append(issues, issue)
 		}
-		if wrapping {
+		if useWrapping {
 			switch {
 			case !o.GetOperation().GetReady():
 				return xerrors.WithStackTrace(errOperationNotReady)
@@ -386,9 +386,9 @@ func (c *conn) NewStream(
 			c.endpoint.Copy(),
 			trace.Method(method),
 		)
-		wrapping = needWrapping(ctx)
-		cc       *grpc.ClientConn
-		s        grpc.ClientStream
+		useWrapping = UseWrapping(ctx)
+		cc          *grpc.ClientConn
+		s           grpc.ClientStream
 	)
 
 	defer func() {
@@ -418,7 +418,7 @@ func (c *conn) NewStream(
 
 	s, err = cc.NewStream(ctx, desc, method, opts...)
 	if err != nil {
-		if wrapping {
+		if useWrapping {
 			err = xerrors.Retryable(
 				xerrors.FromGRPCError(err,
 					xerrors.WithAddress(c.Address()),
@@ -436,7 +436,7 @@ func (c *conn) NewStream(
 	return &grpcClientStream{
 		ClientStream: s,
 		c:            c,
-		wrapping:     wrapping,
+		wrapping:     useWrapping,
 		sentMark:     sentMark,
 		onDone: func(ctx context.Context, md metadata.MD) {
 			cancel()

internal/conn/context.go

@@ -8,7 +8,7 @@ func WithoutWrapping(ctx context.Context) context.Context {
 	return context.WithValue(ctx, ctxNoWrappingKey{}, true)
 }
 
-func needWrapping(ctx context.Context) bool {
+func UseWrapping(ctx context.Context) bool {
 	b, ok := ctx.Value(ctxNoWrappingKey{}).(bool)
 	return !ok || !b
 }

internal/topic/topicreaderinternal/committer.go

@@ -208,7 +208,7 @@ func (c *committer) waitCommitAck(ctx context.Context, waiter commitWaiter) erro
 	case <-ctx.Done():
 		return ctx.Err()
 	case <-waiter.Session.Context().Done():
-		return waiter.Session.Context().Err()
+		return PublicErrCommitSessionToExpiredSession
 	case <-waiter.Committed:
 		return nil
 	}

internal/topic/topicreaderinternal/committer_test.go

@@ -189,7 +189,7 @@ func TestCommitterCommitSync(t *testing.T) {
 		sessionCancel(testErr)
 
 		commitErr := <-waitErr
-		require.ErrorIs(t, commitErr, testErr)
+		require.ErrorIs(t, commitErr, PublicErrCommitSessionToExpiredSession)
 	})
 }
 

internal/topic/topicreaderinternal/partition_session.go

@@ -21,6 +21,9 @@ type partitionSession struct {
 	Topic       string
 	PartitionID int64
 
+	readerID     int64
+	connectionID string
+
 	ctx                context.Context
 	ctxCancel          xcontext.CancelErrFunc
 	partitionSessionID rawtopicreader.PartitionSessionID
@@ -33,6 +36,8 @@ func newPartitionSession(
 	partitionContext context.Context,
 	topic string,
 	partitionID int64,
+	readerID int64,
+	connectionID string,
 	partitionSessionID rawtopicreader.PartitionSessionID,
 	committedOffset rawtopicreader.Offset,
 ) *partitionSession {
@@ -41,6 +46,8 @@ func newPartitionSession(
 	return &partitionSession{
 		Topic:                    topic,
 		PartitionID:              partitionID,
+		readerID:                 readerID,
+		connectionID:             connectionID,
 		ctx:                      partitionContext,
 		ctxCancel:                cancel,
 		partitionSessionID:       partitionSessionID,

internal/topic/topicreaderinternal/reader.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/credentials"
@@ -16,10 +17,19 @@ import (
 )
 
 var (
-	errUnconnected  = xerrors.Retryable(xerrors.Wrap(errors.New("ydb: first connection attempt not finished")))
-	ErrReaderClosed = xerrors.Wrap(errors.New("ydb: reader closed"))
+	errUnconnected = xerrors.Retryable(xerrors.Wrap(
+		errors.New("ydb: first connection attempt not finished"),
+	))
+	errReaderClosed                 = xerrors.Wrap(errors.New("ydb: reader closed"))
+	errCommitSessionFromOtherReader = xerrors.Wrap(errors.New("ydb: commit with session from other reader"))
 )
 
+var globalReaderCounter int64
+
+func nextReaderID() int64 {
+	return atomic.AddInt64(&globalReaderCounter, 1)
+}
+
 //nolint:lll
 //go:generate mockgen -destination raw_topic_reader_stream_mock_test.go -package topicreaderinternal -write_package_comment=false . RawTopicReaderStream
 
@@ -37,6 +47,7 @@ type Reader struct {
 	reader             batchedStreamReader
 	defaultBatchConfig ReadMessageBatchOptions
 	tracer             trace.Topic
+	readerID           int64
 }
 
 type ReadMessageBatchOptions struct {
@@ -78,17 +89,20 @@ func NewReader(
 	opts ...PublicReaderOption,
 ) Reader {
 	cfg := convertNewParamsToStreamConfig(consumer, readSelectors, opts...)
+	readerID := nextReaderID()
+
 	readerConnector := func(ctx context.Context) (batchedStreamReader, error) {
 		stream, err := connector(ctx)
 		if err != nil {
 			return nil, err
 		}
 
-		return newTopicStreamReader(stream, cfg.topicStreamReaderConfig)
+		return newTopicStreamReader(readerID, stream, cfg.topicStreamReaderConfig)
 	}
 
 	res := Reader{
 		reader: newReaderReconnector(
+			readerID,
 			readerConnector,
 			cfg.OperationTimeout(),
 			cfg.RetrySettings,
@@ -97,13 +111,14 @@ func NewReader(
 		),
 		defaultBatchConfig: cfg.DefaultBatchConfig,
 		tracer:             cfg.Tracer,
+		readerID:           readerID,
 	}
 
 	return res
 }
 
 func (r *Reader) Close(ctx context.Context) error {
-	return r.reader.CloseWithError(ctx, xerrors.WithStackTrace(ErrReaderClosed))
+	return r.reader.CloseWithError(ctx, xerrors.WithStackTrace(errReaderClosed))
 }
 
 type readExplicitMessagesCount int
@@ -148,10 +163,21 @@ forReadBatch:
 }
 
 func (r *Reader) Commit(ctx context.Context, offsets PublicCommitRangeGetter) (err error) {
-	return r.reader.Commit(ctx, offsets.getCommitRange().priv)
+	cr := offsets.getCommitRange().priv
+	if cr.partitionSession.readerID != r.readerID {
+		return errCommitSessionFromOtherReader
+	}
+
+	return r.reader.Commit(ctx, cr)
 }
 
 func (r *Reader) CommitRanges(ctx context.Context, ranges []PublicCommitRange) error {
+	for i := range ranges {
+		if ranges[i].priv.partitionSession.readerID != r.readerID {
+			return errCommitSessionFromOtherReader
+		}
+	}
+
 	commitRanges := NewCommitRangesFromPublicCommits(ranges)
 	commitRanges.optimize()
 

internal/topic/topicreaderinternal/reader_test.go

@@ -69,7 +69,11 @@ func TestReader_Close(t *testing.T) {
 		readerReadMessageBatchState := newCallState()
 
 		go func() {
-			readerCommitState.err = reader.Commit(context.Background(), &PublicMessage{})
+			readerCommitState.err = reader.Commit(context.Background(), &PublicMessage{
+				commitRange: commitRange{
+					partitionSession: &partitionSession{},
+				},
+			})
 			close(readerCommitState.callCompleted)
 		}()
 
@@ -100,27 +104,51 @@ func TestReader_Close(t *testing.T) {
 }
 
 func TestReader_Commit(t *testing.T) {
-	mc := gomock.NewController(t)
-	defer mc.Finish()
-
-	baseReader := NewMockbatchedStreamReader(mc)
-	reader := &Reader{reader: baseReader}
-
-	expectedRangeOk := commitRange{
-		commitOffsetStart: 1,
-		commitOffsetEnd:   10,
-		partitionSession:  &partitionSession{partitionSessionID: 10},
-	}
-	baseReader.EXPECT().Commit(gomock.Any(), expectedRangeOk).Return(nil)
-	require.NoError(t, reader.Commit(context.Background(), &PublicMessage{commitRange: expectedRangeOk}))
-
-	expectedRangeErr := commitRange{
-		commitOffsetStart: 15,
-		commitOffsetEnd:   20,
-		partitionSession:  &partitionSession{partitionSessionID: 30},
-	}
-
-	testErr := errors.New("test err")
-	baseReader.EXPECT().Commit(gomock.Any(), expectedRangeErr).Return(testErr)
-	require.ErrorIs(t, reader.Commit(context.Background(), &PublicMessage{commitRange: expectedRangeErr}), testErr)
+	t.Run("OK", func(t *testing.T) {
+		mc := gomock.NewController(t)
+		defer mc.Finish()
+
+		readerID := nextReaderID()
+		baseReader := NewMockbatchedStreamReader(mc)
+		reader := &Reader{
+			reader:   baseReader,
+			readerID: readerID,
+		}
+
+		expectedRangeOk := commitRange{
+			commitOffsetStart: 1,
+			commitOffsetEnd:   10,
+			partitionSession: &partitionSession{
+				readerID:           readerID,
+				partitionSessionID: 10,
+			},
+		}
+		baseReader.EXPECT().Commit(gomock.Any(), expectedRangeOk).Return(nil)
+		require.NoError(t, reader.Commit(context.Background(), &PublicMessage{commitRange: expectedRangeOk}))
+
+		expectedRangeErr := commitRange{
+			commitOffsetStart: 15,
+			commitOffsetEnd:   20,
+			partitionSession: &partitionSession{
+				readerID:           readerID,
+				partitionSessionID: 30,
+			},
+		}
+
+		testErr := errors.New("test err")
+		baseReader.EXPECT().Commit(gomock.Any(), expectedRangeErr).Return(testErr)
+		require.ErrorIs(t, reader.Commit(context.Background(), &PublicMessage{commitRange: expectedRangeErr}), testErr)
+	})
+
+	t.Run("CommitFromOtherReader", func(t *testing.T) {
+		ctx := xtest.Context(t)
+		reader := &Reader{readerID: 1}
+		forCommit := commitRange{
+			commitOffsetStart: 1,
+			commitOffsetEnd:   2,
+			partitionSession:  &partitionSession{readerID: 2},
+		}
+		err := reader.Commit(ctx, forCommit)
+		require.ErrorIs(t, err, errCommitSessionFromOtherReader)
+	})
 }