Fixed topic traces (#1873)

Author: kprokopenko

CHANGELOG.md

@@ -1,4 +1,5 @@
-Fixed race in integration test `TestTopicWriterLogMessagesWithoutData` 
+* Fixed traces handling in `topic.Reader`
+* Fixed race in integration test `TestTopicWriterLogMessagesWithoutData` 
 
 ## v3.117.0
 * Fixed `conn/pool.Get()` behaviour for YDB databases with public IPs. Bug was introduced in v3.116.2

internal/topic/topicreaderinternal/stream_reader_impl.go

@@ -198,16 +198,17 @@ func (r *topicStreamReaderImpl) PopMessagesBatchTx(
 	tx tx.Transaction,
 	opts ReadMessageBatchOptions,
 ) (_ *topicreadercommon.PublicBatch, resErr error) {
+	logCtx := r.cfg.BaseContext
+	onDone := trace.TopicOnReaderStreamPopBatchTx(
+		r.cfg.Trace,
+		&logCtx,
+		r.readerID,
+		r.readConnectionID,
+		tx.SessionID(),
+		tx,
+	)
 	defer func() {
-		logCtx := r.cfg.BaseContext
-		trace.TopicOnReaderStreamPopBatchTx(
-			r.cfg.Trace,
-			&logCtx,
-			r.readerID,
-			r.readConnectionID,
-			tx.SessionID(),
-			tx,
-		)(resErr)
+		onDone(resErr)
 	}()
 
 	batch, err := r.ReadMessageBatch(ctx, opts)
@@ -258,16 +259,17 @@ func (r *topicStreamReaderImpl) commitWithTransaction(
 		r.addOnTransactionCompletedHandler(ctx, tx, batch)
 	} else {
 		_ = retry.Retry(ctx, func(ctx context.Context) (err error) {
+			logCtx := r.cfg.BaseContext
+			onDone := trace.TopicOnReaderTransactionRollback(
+				r.cfg.Trace,
+				&logCtx,
+				r.readerID,
+				r.readConnectionID,
+				tx.SessionID(),
+				tx,
+			)
 			defer func() {
-				logCtx := r.cfg.BaseContext
-				trace.TopicOnReaderTransactionRollback(
-					r.cfg.Trace,
-					&logCtx,
-					r.readerID,
-					r.readConnectionID,
-					tx.SessionID(),
-					tx,
-				)(err)
+				onDone(err)
 			}()
 
 			return tx.Rollback(ctx)
@@ -290,18 +292,17 @@ func (r *topicStreamReaderImpl) addOnTransactionCompletedHandler(
 ) {
 	commitRange := topicreadercommon.GetCommitRange(batch)
 	tx.OnCompleted(func(transactionResult error) {
-		defer func() {
-			logCtx := r.cfg.BaseContext
-			trace.TopicOnReaderTransactionCompleted(
-				r.cfg.Trace,
-				&logCtx,
-				r.readerID,
-				r.readConnectionID,
-				tx.SessionID(),
-				tx,
-				transactionResult,
-			)
-		}()
+		logCtx := r.cfg.BaseContext
+		onDone := trace.TopicOnReaderTransactionCompleted(
+			r.cfg.Trace,
+			&logCtx,
+			r.readerID,
+			r.readConnectionID,
+			tx.SessionID(),
+			tx,
+			transactionResult,
+		)
+		defer onDone()
 
 		if transactionResult == nil {
 			topicreadercommon.BatchGetPartitionSession(batch).SetCommittedOffsetForward(commitRange.CommitOffsetEnd)
@@ -350,38 +351,28 @@ func (r *topicStreamReaderImpl) ReadMessageBatch(
 	ctx context.Context,
 	opts ReadMessageBatchOptions,
 ) (batch *topicreadercommon.PublicBatch, err error) {
+	mergeCtx := xcontext.MergeContexts(ctx, r.cfg.BaseContext)
+	onDone := trace.TopicOnReaderReadMessages(
+		r.cfg.Trace,
+		&mergeCtx,
+		opts.MinCount,
+		opts.MaxCount,
+		r.getRestBufferBytes(),
+	)
 	defer func() {
-		traceFunc := func(
-			messagesCount int,
-			topic string,
-			partitionID int64,
-			partitionSessionID int64,
-			offsetStart int64,
-			offsetEnd int64,
-			freeBufferCapacity int,
-		) {
-			mergeCtx := xcontext.MergeContexts(ctx, r.cfg.BaseContext)
-			onDone := trace.TopicOnReaderReadMessages(
-				r.cfg.Trace,
-				&mergeCtx,
-				opts.MinCount,
-				opts.MaxCount,
-				r.getRestBufferBytes(),
-			)
-			onDone(messagesCount, topic, partitionID, partitionSessionID, offsetStart, offsetEnd, freeBufferCapacity, nil)
-		}
 		if batch == nil {
-			traceFunc(0, "", -1, -1, -1, -1, r.getRestBufferBytes())
+			onDone(0, "", -1, -1, -1, -1, r.getRestBufferBytes(), err)
 		} else {
 			commitRange := topicreadercommon.GetCommitRange(batch)
-			traceFunc(
+			onDone(
 				len(batch.Messages),
 				batch.Topic(),
 				batch.PartitionID(),
 				topicreadercommon.BatchGetPartitionSession(batch).StreamPartitionSessionID.ToInt64(),
 				commitRange.CommitOffsetStart.ToInt64(),
 				commitRange.CommitOffsetEnd.ToInt64(),
 				r.getRestBufferBytes(),
+				err,
 			)
 		}
 	}()
@@ -541,18 +532,20 @@ func (r *topicStreamReaderImpl) Commit(ctx context.Context, commitRange topicrea
 		return xerrors.WithStackTrace(errCommitWithNilPartitionSession)
 	}
 
+	session := commitRange.PartitionSession
+	mergeCtx := xcontext.MergeContexts(ctx, r.cfg.BaseContext)
+	onDone := trace.TopicOnReaderCommit(
+		r.cfg.Trace,
+		&mergeCtx,
+		session.Topic,
+		session.PartitionID,
+		session.StreamPartitionSessionID.ToInt64(),
+		commitRange.CommitOffsetStart.ToInt64(),
+		commitRange.CommitOffsetEnd.ToInt64(),
+	)
+
 	defer func() {
-		session := commitRange.PartitionSession
-		mergeCtx := xcontext.MergeContexts(ctx, r.cfg.BaseContext)
-		trace.TopicOnReaderCommit(
-			r.cfg.Trace,
-			&mergeCtx,
-			session.Topic,
-			session.PartitionID,
-			session.StreamPartitionSessionID.ToInt64(),
-			commitRange.CommitOffsetStart.ToInt64(),
-			commitRange.CommitOffsetEnd.ToInt64(),
-		)
+		onDone(err)
 	}()
 
 	if err = r.checkCommitRange(commitRange); err != nil {
@@ -631,9 +624,10 @@ func (r *topicStreamReaderImpl) setStarted() error {
 func (r *topicStreamReaderImpl) initSession() (err error) {
 	initMessage := topicreadercommon.CreateInitMessage(r.cfg.Consumer, r.cfg.EnableSplitMergeSupport, r.cfg.ReadSelectors)
 
+	logCtx := r.cfg.BaseContext
+	onDone := trace.TopicOnReaderInit(r.cfg.Trace, &logCtx, r.readConnectionID, initMessage)
+
 	defer func() {
-		logCtx := r.cfg.BaseContext
-		onDone := trace.TopicOnReaderInit(r.cfg.Trace, &logCtx, r.readConnectionID, initMessage)
 		onDone(r.readConnectionID, err)
 	}()
 
@@ -822,9 +816,11 @@ func (r *topicStreamReaderImpl) updateTokenLoop(ctx context.Context) {
 
 func (r *topicStreamReaderImpl) onReadResponse(msg *rawtopicreader.ReadResponse) (err error) {
 	resCapacity := r.addRestBufferBytes(-msg.BytesSize)
+
+	logCtx := r.cfg.BaseContext
+	onDone := trace.TopicOnReaderReceiveDataResponse(r.cfg.Trace, &logCtx, r.readConnectionID, resCapacity, msg)
 	defer func() {
-		logCtx := r.cfg.BaseContext
-		trace.TopicOnReaderReceiveDataResponse(r.cfg.Trace, &logCtx, r.readConnectionID, resCapacity, msg)
+		onDone(err)
 	}()
 
 	batches, err2 := topicreadercommon.ReadRawBatchesToPublicBatches(msg, &r.sessionController, r.cfg.Decoders)
@@ -842,9 +838,10 @@ func (r *topicStreamReaderImpl) onReadResponse(msg *rawtopicreader.ReadResponse)
 }
 
 func (r *topicStreamReaderImpl) CloseWithError(ctx context.Context, reason error) (closeErr error) {
+	logCtx := r.cfg.BaseContext
+	onDone := trace.TopicOnReaderClose(r.cfg.Trace, &logCtx, r.readConnectionID, reason)
 	defer func() {
-		logCtx := r.cfg.BaseContext
-		trace.TopicOnReaderClose(r.cfg.Trace, &logCtx, r.readConnectionID, reason)
+		onDone(closeErr)
 	}()
 
 	isFirstClose := false

internal/topic/topicreaderinternal/stream_reader_impl_test.go

@@ -1271,17 +1271,6 @@ func (e *streamEnv) Start() {
 	})
 }
 
-func (e *streamEnv) readerReceiveWaitClose(callback func()) {
-	e.stream.EXPECT().Recv().DoAndReturn(func() (rawtopicreader.ServerMessage, error) {
-		if callback != nil {
-			callback()
-		}
-		<-e.ctx.Done()
-
-		return nil, errors.New("test reader closed")
-	})
-}
-
 func (e *streamEnv) SendFromServer(msg rawtopicreader.ServerMessage) {
 	e.SendFromServerAndSetNextCallback(msg, nil)
 }