add context to topic logs

Author: qrort

internal/grpcwrapper/rawtopic/client.go

@@ -104,6 +104,7 @@ func (c *Client) StreamRead(ctxStreamLifeTime context.Context) (rawtopicreader.S
 func (c *Client) StreamWrite(
 	ctxStreamLifeTime context.Context,
 	tracer *trace.Topic,
+	logContext *context.Context,
 ) (*rawtopicwriter.StreamWriter, error) {
 	protoResp, err := c.service.StreamWrite(ctxStreamLifeTime)
 	if err != nil {
@@ -118,6 +119,7 @@ func (c *Client) StreamWrite(
 		Stream:           protoResp,
 		Tracer:           tracer,
 		InternalStreamID: uuid.New().String(),
+		LogContext:       logContext,
 	}, nil
 }
 

internal/grpcwrapper/rawtopic/rawtopicwriter/streamwriter.go

@@ -1,6 +1,7 @@
 package rawtopicwriter
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"reflect"
@@ -36,6 +37,7 @@ type StreamWriter struct {
 	readMessagesCount    int
 	writtenMessagesCount int
 	sessionID            string
+	LogContext           *context.Context
 }
 
 //nolint:funlen
@@ -52,7 +54,7 @@ func (w *StreamWriter) Recv() (ServerMessage, error) {
 	defer func() {
 		// defer needs for set good session id on first init response before trace the message
 		trace.TopicOnWriterReceiveGRPCMessage(
-			w.Tracer, w.InternalStreamID, w.sessionID, w.readMessagesCount, grpcMsg, sendErr,
+			w.Tracer, w.LogContext, w.InternalStreamID, w.sessionID, w.readMessagesCount, grpcMsg, sendErr,
 		)
 	}()
 	if sendErr != nil {
@@ -141,7 +143,15 @@ func (w *StreamWriter) Send(rawMsg ClientMessage) (err error) {
 
 	err = w.Stream.Send(&protoMsg)
 	w.writtenMessagesCount++
-	trace.TopicOnWriterSentGRPCMessage(w.Tracer, w.InternalStreamID, w.sessionID, w.writtenMessagesCount, &protoMsg, err)
+	trace.TopicOnWriterSentGRPCMessage(
+		w.Tracer,
+		w.LogContext,
+		w.InternalStreamID,
+		w.sessionID,
+		w.writtenMessagesCount,
+		&protoMsg,
+		err,
+	)
 	if err != nil {
 		return xerrors.WithStackTrace(xerrors.Wrap(fmt.Errorf("ydb: failed to send grpc message to writer stream: %w", err)))
 	}

internal/topic/topicclientinternal/client.go

@@ -315,7 +315,10 @@ func (c *Client) StartReader(
 	if err != nil {
 		return nil, err
 	}
-	trace.TopicOnReaderStart(internalReader.Tracer(), internalReader.ID(), consumer, err)
+
+	logCtx := internalReader.GetLogContext()
+	trace.TopicOnReaderStart(internalReader.Tracer(), &logCtx, internalReader.ID(), consumer, err)
+	internalReader.SetLogContext(logCtx)
 
 	return topicreader.NewReader(internalReader), nil
 }
@@ -356,15 +359,16 @@ func (c *Client) createWriterConfig(
 	topicPath string,
 	opts []topicoptions.WriterOption,
 ) topicwriterinternal.WriterReconnectorConfig {
-	var connector topicwriterinternal.ConnectFunc = func(ctx context.Context, tracer *trace.Topic) (
-		topicwriterinternal.RawTopicWriterStream,
-		error,
-	) {
-		return c.rawClient.StreamWrite(ctx, tracer)
-	}
+	//var connector topicwriterinternal.ConnectFunc = func(ctx context.Context, tracer *trace.Topic) (
+	//	topicwriterinternal.RawTopicWriterStream,
+	//	error,
+	//) {
+	//	return c.rawClient.StreamWrite(ctx, tracer)
+	//}
 
 	options := []topicoptions.WriterOption{
-		topicwriterinternal.WithConnectFunc(connector),
+		topicwriterinternal.WithRawClient(&c.rawClient),
+		// topicwriterinternal.WithConnectFunc(connector),
 		topicwriterinternal.WithTopic(topicPath),
 		topicwriterinternal.WithCommonConfig(c.cfg.Common),
 		topicwriterinternal.WithTrace(c.cfg.Trace),

internal/topic/topiclistenerinternal/stream_listener.go

@@ -141,7 +141,6 @@ func (l *streamListener) initVars(sessionIDCounter *atomic.Int64) {
 	}
 }
 
-//nolint:funlen
 func (l *streamListener) initStream(ctx context.Context, client TopicClient) error {
 	streamCtx, streamClose := context.WithCancelCause(xcontext.ValueOnly(ctx))
 	l.streamClose = streamClose

internal/topic/topicreadercommon/committer.go

@@ -148,6 +148,7 @@ func (c *Committer) pushCommitsLoop(ctx context.Context) {
 
 		onDone := trace.TopicOnReaderSendCommitMessage(
 			c.tracer,
+			&ctx,
 			&commits,
 		)
 		err := c.send(commits.ToRawMessage())

internal/topic/topicreaderinternal/batched_stream_reader_interface.go

@@ -15,4 +15,6 @@ type batchedStreamReader interface {
 	Commit(ctx context.Context, commitRange topicreadercommon.CommitRange) error
 	CloseWithError(ctx context.Context, err error) error
 	PopMessagesBatchTx(ctx context.Context, tx tx.Transaction, opts ReadMessageBatchOptions) (*topicreadercommon.PublicBatch, error) //nolint:lll
+	GetLogContext() context.Context
+	SetLogContext(ctx context.Context)
 }

internal/topic/topicreaderinternal/batched_stream_reader_mock_test.go

@@ -35,6 +35,13 @@ type Reader struct {
 	readerID           int64
 }
 
+func (r *Reader) GetLogContext() context.Context {
+	return r.reader.GetLogContext()
+}
+func (r *Reader) SetLogContext(ctx context.Context) {
+	r.reader.SetLogContext(ctx)
+}
+
 type ReadMessageBatchOptions struct {
 	batcherGetOptions
 }
@@ -89,14 +96,17 @@ func NewReader(
 		return newTopicStreamReader(client, readerID, stream, cfg.topicStreamReaderConfig)
 	}
 
+	reader := newReaderReconnector(
+		readerID,
+		readerConnector,
+		cfg.OperationTimeout(),
+		cfg.BaseContext,
+		cfg.RetrySettings,
+		cfg.Trace,
+	)
+
 	res := Reader{
-		reader: newReaderReconnector(
-			readerID,
-			readerConnector,
-			cfg.OperationTimeout(),
-			cfg.RetrySettings,
-			cfg.Trace,
-		),
+		reader:             reader,
 		defaultBatchConfig: cfg.DefaultBatchConfig,
 		tracer:             cfg.Trace,
 		readerID:           readerID,