Add context to topic logs

Author: qrort

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added `WithReaderLogContext`, `WithWriterLogContext` options to topic reader/writer to supply log entries with user context fields
+
 ## v3.108.3
 * Fixed handling of zero values for DyNumber
 * Fixed the decimal yql slice bounds out of range

go.sum

@@ -11,7 +11,6 @@ github.com/cncf/udpa/go v0.0.0-20210930031921-04548b0d99d4/go.mod h1:6pvJx4me5XP
 github.com/cncf/xds/go v0.0.0-20210922020428-25de7278fc84/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
 github.com/cncf/xds/go v0.0.0-20211001041855-01bcc9b48dfe/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
 github.com/cncf/xds/go v0.0.0-20211011173535-cb28da3451f1/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
-github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -72,11 +71,13 @@ github.com/rekby/fixenv v0.6.1/go.mod h1:/b5LRc06BYJtslRtHKxsPWFT/ySpHV+rWvzTg+X
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
+github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
+github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
+github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/ydb-platform/ydb-go-genproto v0.0.0-20241112172322-ea1f63298f77 h1:LY6cI8cP4B9rrpTleZk95+08kl2gF4rixG7+V/dwL6Q=
 github.com/ydb-platform/ydb-go-genproto v0.0.0-20241112172322-ea1f63298f77/go.mod h1:Er+FePu1dNUieD+XTMDduGpQuCPssK5Q4BjF+IIXJ3I=
 go.opentelemetry.io/otel v1.31.0 h1:NsJcKPIW0D0H3NgzPDHmo0WW6SptzPdqg/L1zsIm2hY=
@@ -94,6 +95,10 @@ go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
 go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=
 go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc=
+go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
+go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
+go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8=
+go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=

internal/grpcwrapper/rawtopic/client.go

@@ -126,6 +126,7 @@ func (c *Client) StreamWrite(
 		Stream:           protoResp,
 		Tracer:           tracer,
 		InternalStreamID: uuid.New().String(),
+		LogContext:       &ctxStreamLifeTime,
 	}, nil
 }
 

internal/grpcwrapper/rawtopic/rawtopicwriter/streamwriter.go

@@ -1,6 +1,7 @@
 package rawtopicwriter
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"reflect"
@@ -34,6 +35,7 @@ type StreamWriter struct {
 	readMessagesCount    int
 	writtenMessagesCount int
 	sessionID            string
+	LogContext           *context.Context
 }
 
 //nolint:funlen
@@ -50,7 +52,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 {
@@ -139,7 +141,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

@@ -322,7 +322,8 @@ func (c *Client) StartReader(
 	if err != nil {
 		return nil, err
 	}
-	trace.TopicOnReaderStart(internalReader.Tracer(), internalReader.ID(), consumer, err)
+
+	internalReader.TopicOnReaderStart(consumer, err)
 
 	return topicreader.NewReader(internalReader), nil
 }
@@ -365,15 +366,8 @@ 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)
-	}
-
 	options := []topicoptions.WriterOption{
-		topicwriterinternal.WithConnectFunc(connector),
+		topicwriterinternal.WithRawClient(&c.rawClient),
 		topicwriterinternal.WithTopic(topicPath),
 		topicwriterinternal.WithCommonConfig(c.cfg.Common),
 		topicwriterinternal.WithTrace(c.cfg.Trace),

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,5 @@ 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
+	TopicOnReaderStart(consumer string, err error)
 }

internal/topic/topicreaderinternal/batched_stream_reader_mock_test.go

@@ -39,6 +39,10 @@ type Reader struct {
 	readerID           int64
 }
 
+func (r *Reader) TopicOnReaderStart(consumer string, err error) {
+	r.reader.TopicOnReaderStart(consumer, err)
+}
+
 type ReadMessageBatchOptions struct {
 	batcherGetOptions
 }
@@ -93,14 +97,17 @@ func NewReader(
 		return newTopicStreamReader(client, readerID, stream, cfg.topicStreamReaderConfig)
 	}
 
+	reader := newReaderReconnector(
+		cfg.BaseContext,
+		readerID,
+		readerConnector,
+		cfg.OperationTimeout(),
+		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,