Merge pull request #519 separate commit for expired session and other reader

rekby · web-flow · commit 79d9a27978b9 · 2023-01-13T12:26:29.000+03:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,5 @@
+* Separated errors of commit from other reader and to expired session
+
 ## v3.42.4
 * Added `ydb.WithDisableServerBalancer()` database/sql connector option
 
diff --git a/internal/topic/topicreaderinternal/committer.go 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
 	}
diff --git a/internal/topic/topicreaderinternal/committer_test.go b/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)
 	})
 }
 
diff --git a/internal/topic/topicreaderinternal/partition_session.go b/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,
diff --git a/internal/topic/topicreaderinternal/reader.go b/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()
 
diff --git a/internal/topic/topicreaderinternal/reader_test.go b/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)
+	})
 }
diff --git a/internal/topic/topicreaderinternal/stream_reader_impl.go b/internal/topic/topicreaderinternal/stream_reader_impl.go
@@ -23,9 +23,10 @@ import (
 )
 
 var (
+	PublicErrCommitSessionToExpiredSession = xerrors.Wrap(errors.New("ydb: commit to expired session"))
+
 	errPartitionSessionStoppedByServer = xerrors.Wrap(errors.New("ydb: topic partition session stopped by server"))
 	errPartitionSessionStoppedBySDK    = xerrors.Wrap(errors.New("ydb: topic partition session stopped by sdk"))
-	errCommitSessionFromOtherReader    = xerrors.Wrap(errors.New("ydb: commit with session from other reader"))
 	errCommitWithNilPartitionSession   = xerrors.Wrap(errors.New("ydb: commit with nil partition session"))
 )
 
@@ -48,6 +49,7 @@ type topicStreamReaderImpl struct {
 
 	stream           RawTopicReaderStream
 	readConnectionID string
+	readerID         int64
 
 	m       xsync.RWMutex
 	err     error
@@ -106,6 +108,7 @@ func (cfg *topicStreamReaderConfig) initMessage() *rawtopicreader.InitRequest {
 }
 
 func newTopicStreamReader(
+	readerID int64,
 	stream RawTopicReaderStream,
 	cfg topicStreamReaderConfig,
 ) (_ *topicStreamReaderImpl, err error) {
@@ -115,7 +118,7 @@ func newTopicStreamReader(
 		}
 	}()
 
-	reader := newTopicStreamReaderStopped(stream, cfg)
+	reader := newTopicStreamReaderStopped(readerID, stream, cfg)
 	if err = reader.initSession(); err != nil {
 		return nil, err
 	}
@@ -127,6 +130,7 @@ func newTopicStreamReader(
 }
 
 func newTopicStreamReaderStopped(
+	readerID int64,
 	stream RawTopicReaderStream,
 	cfg topicStreamReaderConfig,
 ) *topicStreamReaderImpl {
@@ -147,6 +151,7 @@ func newTopicStreamReaderStopped(
 		batcher:               newBatcher(),
 		backgroundWorkers:     *background.NewWorker(stopPump),
 		readConnectionID:      "preinitID-" + readerConnectionID.String(),
+		readerID:              readerID,
 		rawMessagesFromBuffer: make(chan rawtopicreader.ServerMessage, 1),
 	}
 
@@ -361,7 +366,7 @@ func (r *topicStreamReaderImpl) checkCommitRange(commitRange commitRange) error
 
 	ownSession, err := r.sessionController.Get(session.partitionSessionID)
 	if err != nil || session != ownSession {
-		return xerrors.WithStackTrace(errCommitSessionFromOtherReader)
+		return xerrors.WithStackTrace(PublicErrCommitSessionToExpiredSession)
 	}
 
 	return nil
@@ -706,6 +711,8 @@ func (r *topicStreamReaderImpl) onStartPartitionSessionRequest(m *rawtopicreader
 		r.ctx,
 		m.PartitionSession.Path,
 		m.PartitionSession.PartitionID,
+		r.readerID,
+		r.readConnectionID,
 		m.PartitionSession.PartitionSessionID,
 		m.CommittedOffset,
 	)
diff --git a/internal/topic/topicreaderinternal/stream_reader_impl_test.go b/internal/topic/topicreaderinternal/stream_reader_impl_test.go
@@ -197,7 +197,7 @@ func TestTopicStreamReaderImpl_Create(t *testing.T) {
 		}, nil)
 		stream.EXPECT().CloseSend().Return(nil)
 
-		reader, err := newTopicStreamReader(stream, newTopicStreamReaderConfig())
+		reader, err := newTopicStreamReader(nextReaderID(), stream, newTopicStreamReaderConfig())
 		require.Error(t, err)
 		require.Nil(t, reader)
 	})
@@ -686,7 +686,15 @@ func TestTopicStreamReadImpl_CommitWithBadSession(t *testing.T) {
 	e.Start()
 
 	cr := commitRange{
-		partitionSession: newPartitionSession(context.Background(), "asd", 123, 222, 213),
+		partitionSession: newPartitionSession(
+			context.Background(),
+			"asd",
+			123,
+			nextReaderID(),
+			"bad-connection-id",
+			222,
+			213,
+		),
 	}
 	err := e.reader.Commit(e.ctx, cr)
 	require.Error(t, err)
@@ -733,14 +741,22 @@ func newTopicReaderTestEnv(t testing.TB) streamEnv {
 	cfg.BufferSizeProtoBytes = initialBufferSizeBytes
 	cfg.CommitterBatchTimeLag = 0
 
-	reader := newTopicStreamReaderStopped(stream, cfg)
+	reader := newTopicStreamReaderStopped(nextReaderID(), stream, cfg)
 	// reader.initSession() - skip stream level initialization
 
 	const testPartitionID = 5
 	const testSessionID = 15
 	const testSessionComitted = 20
 
-	session := newPartitionSession(ctx, "/test", testPartitionID, testSessionID, testSessionComitted)
+	session := newPartitionSession(
+		ctx,
+		"/test",
+		testPartitionID,
+		reader.readerID,
+		reader.readConnectionID,
+		testSessionID,
+		testSessionComitted,
+	)
 	require.NoError(t, reader.sessionController.Add(session))
 
 	env := streamEnv{
diff --git a/internal/topic/topicreaderinternal/stream_reconnector.go b/internal/topic/topicreaderinternal/stream_reconnector.go
@@ -42,6 +42,7 @@ type readerReconnector struct {
 	connectTimeout         time.Duration
 
 	closeOnce sync.Once
+	readerID  int64
 
 	m                          xsync.RWMutex
 	streamConnectionInProgress empty.Chan // opened if connection in progress, closed if connection established
@@ -52,13 +53,15 @@ type readerReconnector struct {
 
 //nolint:revive
 func newReaderReconnector(
+	readerID int64,
 	connector readerConnectFunc,
 	connectTimeout time.Duration,
 	retrySettings topic.RetrySettings,
 	tracer trace.Topic,
 	baseContext context.Context,
 ) *readerReconnector {
 	res := &readerReconnector{
+		readerID:       readerID,
 		clock:          clockwork.NewRealClock(),
 		readerConnect:  connector,
 		streamErr:      errUnconnected,
@@ -138,7 +141,7 @@ func (r *readerReconnector) CloseWithError(ctx context.Context, err error) error
 		closeErr = r.background.Close(ctx, err)
 
 		if r.streamVal != nil {
-			streamCloseErr := r.streamVal.CloseWithError(ctx, xerrors.WithStackTrace(ErrReaderClosed))
+			streamCloseErr := r.streamVal.CloseWithError(ctx, xerrors.WithStackTrace(errReaderClosed))
 			if closeErr == nil {
 				closeErr = streamCloseErr
 			}
diff --git a/internal/topic/topicreaderinternal/stream_reconnector_test.go b/internal/topic/topicreaderinternal/stream_reconnector_test.go
diff --git a/topic/topicreader/errors.go b/topic/topicreader/errors.go
diff --git a/topic/topicreader/reader.go b/topic/topicreader/reader.go

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+* Separated errors of commit from other reader and to expired session`
	`2`	`+`
`1`	`3`	`## v3.42.4`
`2`	`4`	* Added `ydb.WithDisableServerBalancer()` database/sql connector option
`3`	`5`
Original file line number	Diff line number	Diff line change
`@@ -208,7 +208,7 @@ func (c *committer) waitCommitAck(ctx context.Context, waiter commitWaiter) erro`
`208`	`208`	`case <-ctx.Done():`
`209`	`209`	`return ctx.Err()`
`210`	`210`	`case <-waiter.Session.Context().Done():`
`211`		`- return waiter.Session.Context().Err()`
	`211`	`+ return PublicErrCommitSessionToExpiredSession`
`212`	`212`	`case <-waiter.Committed:`
`213`	`213`	`return nil`
`214`	`214`	`}`
Original file line number	Diff line number	Diff line change
`@@ -189,7 +189,7 @@ func TestCommitterCommitSync(t *testing.T) {`
`189`	`189`	`sessionCancel(testErr)`
`190`	`190`
`191`	`191`	`commitErr := <-waitErr`
`192`		`- require.ErrorIs(t, commitErr, testErr)`
	`192`	`+ require.ErrorIs(t, commitErr, PublicErrCommitSessionToExpiredSession)`
`193`	`193`	`})`
`194`	`194`	`}`
`195`	`195`