Skip to content

Commit e5def5a

Browse files
qrortqrort
committed
add context to topic logs
1 parent e5c21a8 commit e5def5a

File tree

21 files changed

+310
-141
lines changed

21 files changed

+310
-141
lines changed

‎internal/grpcwrapper/rawtopic/client.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ func (c *Client) StreamRead(ctxStreamLifeTime context.Context) (rawtopicreader.S
104104
func (c *Client) StreamWrite(
105105
ctxStreamLifeTime context.Context,
106106
tracer *trace.Topic,
107+
logContext *context.Context,
107108
) (*rawtopicwriter.StreamWriter, error) {
108109
protoResp, err := c.service.StreamWrite(ctxStreamLifeTime)
109110
if err != nil {
@@ -118,6 +119,7 @@ func (c *Client) StreamWrite(
118119
Stream: protoResp,
119120
Tracer: tracer,
120121
InternalStreamID: uuid.New().String(),
122+
LogContext: logContext,
121123
}, nil
122124
}
123125

‎internal/grpcwrapper/rawtopic/rawtopicwriter/streamwriter.go

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package rawtopicwriter
22

33
import (
4+
"context"
45
"errors"
56
"fmt"
67
"reflect"
@@ -36,6 +37,7 @@ type StreamWriter struct {
3637
readMessagesCount int
3738
writtenMessagesCount int
3839
sessionID string
40+
LogContext *context.Context
3941
}
4042

4143
//nolint:funlen
@@ -52,7 +54,7 @@ func (w *StreamWriter) Recv() (ServerMessage, error) {
5254
defer func() {
5355
// defer needs for set good session id on first init response before trace the message
5456
trace.TopicOnWriterReceiveGRPCMessage(
55-
w.Tracer, w.InternalStreamID, w.sessionID, w.readMessagesCount, grpcMsg, sendErr,
57+
w.Tracer, w.LogContext, w.InternalStreamID, w.sessionID, w.readMessagesCount, grpcMsg, sendErr,
5658
)
5759
}()
5860
if sendErr != nil {
@@ -141,7 +143,15 @@ func (w *StreamWriter) Send(rawMsg ClientMessage) (err error) {
141143

142144
err = w.Stream.Send(&protoMsg)
143145
w.writtenMessagesCount++
144-
trace.TopicOnWriterSentGRPCMessage(w.Tracer, w.InternalStreamID, w.sessionID, w.writtenMessagesCount, &protoMsg, err)
146+
trace.TopicOnWriterSentGRPCMessage(
147+
w.Tracer,
148+
w.LogContext,
149+
w.InternalStreamID,
150+
w.sessionID,
151+
w.writtenMessagesCount,
152+
&protoMsg,
153+
err,
154+
)
145155
if err != nil {
146156
return xerrors.WithStackTrace(xerrors.Wrap(fmt.Errorf("ydb: failed to send grpc message to writer stream: %w", err)))
147157
}

‎internal/topic/topicclientinternal/client.go

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -311,11 +311,11 @@ func (c *Client) StartReader(
311311
}
312312
opts = append(defaultOpts, opts...)
313313

314-
internalReader, err := topicreaderinternal.NewReader(&c.rawClient, connector, consumer, readSelectors, opts...)
314+
internalReader, logCtx, err := topicreaderinternal.NewReader(&c.rawClient, connector, consumer, readSelectors, opts...)
315315
if err != nil {
316316
return nil, err
317317
}
318-
trace.TopicOnReaderStart(internalReader.Tracer(), internalReader.ID(), consumer, err)
318+
trace.TopicOnReaderStart(internalReader.Tracer(), logCtx, internalReader.ID(), consumer, err)
319319

320320
return topicreader.NewReader(internalReader), nil
321321
}
@@ -356,15 +356,16 @@ func (c *Client) createWriterConfig(
356356
topicPath string,
357357
opts []topicoptions.WriterOption,
358358
) topicwriterinternal.WriterReconnectorConfig {
359-
var connector topicwriterinternal.ConnectFunc = func(ctx context.Context, tracer *trace.Topic) (
360-
topicwriterinternal.RawTopicWriterStream,
361-
error,
362-
) {
363-
return c.rawClient.StreamWrite(ctx, tracer)
364-
}
359+
//var connector topicwriterinternal.ConnectFunc = func(ctx context.Context, tracer *trace.Topic) (
360+
// topicwriterinternal.RawTopicWriterStream,
361+
// error,
362+
//) {
363+
// return c.rawClient.StreamWrite(ctx, tracer)
364+
//}
365365

366366
options := []topicoptions.WriterOption{
367-
topicwriterinternal.WithConnectFunc(connector),
367+
topicwriterinternal.WithRawClient(&c.rawClient),
368+
// topicwriterinternal.WithConnectFunc(connector),
368369
topicwriterinternal.WithTopic(topicPath),
369370
topicwriterinternal.WithCommonConfig(c.cfg.Common),
370371
topicwriterinternal.WithTrace(c.cfg.Trace),

‎internal/topic/topiclistenerinternal/stream_listener.go

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,6 @@ func (l *streamListener) initVars(sessionIDCounter *atomic.Int64) {
141141
}
142142
}
143143

144-
//nolint:funlen
145144
func (l *streamListener) initStream(ctx context.Context, client TopicClient) error {
146145
streamCtx, streamClose := context.WithCancelCause(xcontext.ValueOnly(ctx))
147146
l.streamClose = streamClose

‎internal/topic/topicreadercommon/committer.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,7 @@ func (c *Committer) pushCommitsLoop(ctx context.Context) {
148148

149149
onDone := trace.TopicOnReaderSendCommitMessage(
150150
c.tracer,
151+
&ctx,
151152
&commits,
152153
)
153154
err := c.send(commits.ToRawMessage())

‎internal/topic/topicreaderinternal/reader.go

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -68,11 +68,11 @@ func NewReader(
6868
consumer string,
6969
readSelectors []topicreadercommon.PublicReadSelector,
7070
opts ...PublicReaderOption,
71-
) (Reader, error) {
71+
) (Reader, *context.Context, error) {
7272
cfg := convertNewParamsToStreamConfig(consumer, readSelectors, opts...)
7373

7474
if errs := cfg.Validate(); len(errs) > 0 {
75-
return Reader{}, xerrors.WithStackTrace(fmt.Errorf(
75+
return Reader{}, nil, xerrors.WithStackTrace(fmt.Errorf(
7676
"ydb: failed to start topic reader, because is contains error in config: %w",
7777
errors.Join(errs...),
7878
))
@@ -89,20 +89,23 @@ func NewReader(
8989
return newTopicStreamReader(client, readerID, stream, cfg.topicStreamReaderConfig)
9090
}
9191

92+
reader := newReaderReconnector(
93+
readerID,
94+
readerConnector,
95+
cfg.OperationTimeout(),
96+
&cfg.BaseContext,
97+
cfg.RetrySettings,
98+
cfg.Trace,
99+
)
100+
92101
res := Reader{
93-
reader: newReaderReconnector(
94-
readerID,
95-
readerConnector,
96-
cfg.OperationTimeout(),
97-
cfg.RetrySettings,
98-
cfg.Trace,
99-
),
102+
reader: reader,
100103
defaultBatchConfig: cfg.DefaultBatchConfig,
101104
tracer: cfg.Trace,
102105
readerID: readerID,
103106
}
104107

105-
return res, nil
108+
return res, reader.logContext, nil
106109
}
107110

108111
func (r *Reader) WaitInit(ctx context.Context) error {

‎internal/topic/topicreaderinternal/stream_reader_impl.go

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -574,7 +574,7 @@ func (r *topicStreamReaderImpl) checkCommitRange(commitRange topicreadercommon.C
574574
func (r *topicStreamReaderImpl) send(msg rawtopicreader.ClientMessage) error {
575575
err := r.stream.Send(msg)
576576
if err != nil {
577-
trace.TopicOnReaderError(r.cfg.Trace, r.readConnectionID, err)
577+
trace.TopicOnReaderError(r.cfg.Trace, &r.ctx, r.readConnectionID, err)
578578
_ = r.CloseWithError(r.ctx, err)
579579
}
580580

@@ -613,7 +613,7 @@ func (r *topicStreamReaderImpl) setStarted() error {
613613
func (r *topicStreamReaderImpl) initSession() (err error) {
614614
initMessage := topicreadercommon.CreateInitMessage(r.cfg.Consumer, r.cfg.ReadSelectors)
615615

616-
onDone := trace.TopicOnReaderInit(r.cfg.Trace, r.readConnectionID, initMessage)
616+
onDone := trace.TopicOnReaderInit(r.cfg.Trace, &r.ctx, r.readConnectionID, initMessage)
617617
defer func() {
618618
onDone(r.readConnectionID, err)
619619
}()
@@ -664,9 +664,9 @@ func (r *topicStreamReaderImpl) readMessagesLoop(ctx context.Context) {
664664
for {
665665
serverMessage, err := r.stream.Recv()
666666
if err != nil {
667-
trace.TopicOnReaderError(r.cfg.Trace, r.readConnectionID, err)
667+
trace.TopicOnReaderError(r.cfg.Trace, &ctx, r.readConnectionID, err)
668668
if errors.Is(err, rawtopicreader.ErrUnexpectedMessageType) {
669-
trace.TopicOnReaderUnknownGrpcMessage(r.cfg.Trace, r.readConnectionID, err)
669+
trace.TopicOnReaderUnknownGrpcMessage(r.cfg.Trace, &ctx, r.readConnectionID, err)
670670
// new messages can be added to protocol, it must be backward compatible to old programs
671671
// and skip message is safe
672672
continue
@@ -687,7 +687,7 @@ func (r *topicStreamReaderImpl) readMessagesLoop(ctx context.Context) {
687687

688688
switch m := serverMessage.(type) {
689689
case *rawtopicreader.ReadResponse:
690-
if err = r.onReadResponse(m); err != nil {
690+
if err = r.onReadResponse(ctx, m); err != nil {
691691
_ = r.CloseWithError(ctx, err)
692692
}
693693
case *rawtopicreader.StartPartitionSessionRequest:
@@ -703,7 +703,7 @@ func (r *topicStreamReaderImpl) readMessagesLoop(ctx context.Context) {
703703
return
704704
}
705705
case *rawtopicreader.CommitOffsetResponse:
706-
if err = r.onCommitResponse(m); err != nil {
706+
if err = r.onCommitResponse(ctx, m); err != nil {
707707
_ = r.CloseWithError(ctx, err)
708708

709709
return
@@ -714,6 +714,7 @@ func (r *topicStreamReaderImpl) readMessagesLoop(ctx context.Context) {
714714
default:
715715
trace.TopicOnReaderUnknownGrpcMessage(
716716
r.cfg.Trace,
717+
&ctx,
717718
r.readConnectionID,
718719
xerrors.WithStackTrace(xerrors.Wrap(fmt.Errorf(
719720
"ydb: unexpected message type in stream reader: %v",
@@ -753,7 +754,7 @@ func (r *topicStreamReaderImpl) dataRequestLoop(ctx context.Context) {
753754
}
754755

755756
resCapacity := r.addRestBufferBytes(sum)
756-
trace.TopicOnReaderSentDataRequest(r.cfg.Trace, r.readConnectionID, sum, resCapacity)
757+
trace.TopicOnReaderSentDataRequest(r.cfg.Trace, &ctx, r.readConnectionID, sum, resCapacity)
757758
if err := r.sendDataRequest(sum); err != nil {
758759
return
759760
}
@@ -791,9 +792,9 @@ func (r *topicStreamReaderImpl) updateTokenLoop(ctx context.Context) {
791792
}
792793
}
793794

794-
func (r *topicStreamReaderImpl) onReadResponse(msg *rawtopicreader.ReadResponse) (err error) {
795+
func (r *topicStreamReaderImpl) onReadResponse(ctx context.Context, msg *rawtopicreader.ReadResponse) (err error) {
795796
resCapacity := r.addRestBufferBytes(-msg.BytesSize)
796-
onDone := trace.TopicOnReaderReceiveDataResponse(r.cfg.Trace, r.readConnectionID, resCapacity, msg)
797+
onDone := trace.TopicOnReaderReceiveDataResponse(r.cfg.Trace, &ctx, r.readConnectionID, resCapacity, msg)
797798
defer func() {
798799
onDone(err)
799800
}()
@@ -813,7 +814,7 @@ func (r *topicStreamReaderImpl) onReadResponse(msg *rawtopicreader.ReadResponse)
813814
}
814815

815816
func (r *topicStreamReaderImpl) CloseWithError(ctx context.Context, reason error) (closeErr error) {
816-
onDone := trace.TopicOnReaderClose(r.cfg.Trace, r.readConnectionID, reason)
817+
onDone := trace.TopicOnReaderClose(r.cfg.Trace, &ctx, r.readConnectionID, reason)
817818
defer onDone(closeErr)
818819

819820
isFirstClose := false
@@ -853,7 +854,7 @@ func (r *topicStreamReaderImpl) CloseWithError(ctx context.Context, reason error
853854
return closeErr
854855
}
855856

856-
func (r *topicStreamReaderImpl) onCommitResponse(msg *rawtopicreader.CommitOffsetResponse) error {
857+
func (r *topicStreamReaderImpl) onCommitResponse(ctx context.Context, msg *rawtopicreader.CommitOffsetResponse) error {
857858
for i := range msg.PartitionsCommittedOffsets {
858859
commit := &msg.PartitionsCommittedOffsets[i]
859860
partition, err := r.sessionController.Get(commit.PartitionSessionID)
@@ -864,6 +865,7 @@ func (r *topicStreamReaderImpl) onCommitResponse(msg *rawtopicreader.CommitOffse
864865

865866
trace.TopicOnReaderCommittedNotify(
866867
r.cfg.Trace,
868+
&ctx,
867869
r.readConnectionID,
868870
partition.Topic,
869871
partition.PartitionID,
@@ -880,10 +882,11 @@ func (r *topicStreamReaderImpl) onCommitResponse(msg *rawtopicreader.CommitOffse
880882
func (r *topicStreamReaderImpl) updateToken(ctx context.Context) {
881883
onUpdateToken := trace.TopicOnReaderUpdateToken(
882884
r.cfg.Trace,
885+
&ctx,
883886
r.readConnectionID,
884887
)
885888
token, err := r.cfg.Cred.Token(ctx)
886-
onSent := onUpdateToken(len(token), err)
889+
onSent := onUpdateToken(&ctx, len(token), err)
887890
if err != nil {
888891
return
889892
}

‎internal/topic/topicreaderinternal/stream_reconnector.go

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ type readerConnectFunc func(ctx context.Context) (batchedStreamReader, error)
3434
type readerReconnector struct {
3535
background background.Worker
3636
clock clockwork.Clock
37+
logContext *context.Context
3738
retrySettings topic.RetrySettings
3839
streamVal batchedStreamReader
3940
streamContextCancel context.CancelCauseFunc
@@ -55,6 +56,7 @@ func newReaderReconnector(
5556
readerID int64,
5657
connector readerConnectFunc,
5758
connectTimeout time.Duration,
59+
logContext *context.Context,
5860
retrySettings topic.RetrySettings,
5961
tracer *trace.Topic,
6062
) *readerReconnector {
@@ -64,6 +66,7 @@ func newReaderReconnector(
6466
readerConnect: connector,
6567
streamErr: errUnconnected,
6668
connectTimeout: connectTimeout,
69+
logContext: logContext,
6770
tracer: tracer,
6871
retrySettings: retrySettings,
6972
}
@@ -263,7 +266,7 @@ func (r *readerReconnector) reconnectionLoop(ctx context.Context) {
263266
}
264267
}
265268

266-
onReconnectionDone := trace.TopicOnReaderReconnect(r.tracer, request.reason)
269+
onReconnectionDone := trace.TopicOnReaderReconnect(r.tracer, &ctx, request.reason)
267270

268271
if request.reason != nil {
269272
retryBackoff, stopRetryReason := r.checkErrRetryMode(
@@ -299,7 +302,7 @@ func (r *readerReconnector) reconnectionLoop(ctx context.Context) {
299302

300303
//nolint:funlen
301304
func (r *readerReconnector) reconnect(ctx context.Context, reason error, oldReader batchedStreamReader) (err error) {
302-
onDone := trace.TopicOnReaderReconnect(r.tracer, reason)
305+
onDone := trace.TopicOnReaderReconnect(r.tracer, &ctx, reason)
303306
defer func() { onDone(err) }()
304307

305308
if err = ctx.Err(); err != nil {
@@ -339,9 +342,9 @@ func (r *readerReconnector) reconnect(ctx context.Context, reason error, oldRead
339342
r.background.Start("ydb topic reader send reconnect message", func(ctx context.Context) {
340343
select {
341344
case r.reconnectFromBadStream <- newReconnectRequest(oldReader, sendReason):
342-
trace.TopicOnReaderReconnectRequest(r.tracer, err, true)
345+
trace.TopicOnReaderReconnectRequest(r.tracer, &ctx, err, true)
343346
case <-ctx.Done():
344-
trace.TopicOnReaderReconnectRequest(r.tracer, ctx.Err(), false)
347+
trace.TopicOnReaderReconnectRequest(r.tracer, &ctx, ctx.Err(), false)
345348
}
346349
})
347350
default:
@@ -441,10 +444,10 @@ func (r *readerReconnector) fireReconnectOnRetryableError(stream batchedStreamRe
441444
select {
442445
case r.reconnectFromBadStream <- newReconnectRequest(stream, err):
443446
// send signal
444-
trace.TopicOnReaderReconnectRequest(r.tracer, err, true)
447+
trace.TopicOnReaderReconnectRequest(r.tracer, r.logContext, err, true)
445448
default:
446449
// previous reconnect signal in process, no need sent signal more
447-
trace.TopicOnReaderReconnectRequest(r.tracer, err, false)
450+
trace.TopicOnReaderReconnectRequest(r.tracer, r.logContext, err, false)
448451
}
449452
}
450453

‎internal/topic/topicwriterinternal/encoders.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package topicwriterinternal
33
import (
44
"bytes"
55
"compress/gzip"
6+
"context"
67
"fmt"
78
"io"
89
"sync"
@@ -164,6 +165,7 @@ type EncoderSelector struct {
164165
parallelCompressors int
165166
batchCounter int
166167
measureIntervalBatches int
168+
logContext *context.Context
167169
}
168170

169171
func NewEncoderSelector(
@@ -172,6 +174,7 @@ func NewEncoderSelector(
172174
parallelCompressors int,
173175
tracer *trace.Topic,
174176
writerReconnectorID, sessionID string,
177+
logContext *context.Context,
175178
) EncoderSelector {
176179
if parallelCompressors <= 0 {
177180
panic("ydb: need leas one allowed compressor")
@@ -184,6 +187,7 @@ func NewEncoderSelector(
184187
tracer: tracer,
185188
writerReconnectorID: writerReconnectorID,
186189
sessionID: sessionID,
190+
logContext: logContext,
187191
}
188192
res.ResetAllowedCodecs(allowedCodecs)
189193

@@ -195,6 +199,7 @@ func (s *EncoderSelector) CompressMessages(messages []messageWithDataContent) (r
195199
if err == nil {
196200
onCompressDone := trace.TopicOnWriterCompressMessages(
197201
s.tracer,
202+
s.logContext,
198203
s.writerReconnectorID,
199204
s.sessionID,
200205
codec.ToInt32(),
@@ -265,6 +270,7 @@ func (s *EncoderSelector) measureCodecs(messages []messageWithDataContent) (rawt
265270
}
266271
onCompressDone := trace.TopicOnWriterCompressMessages(
267272
s.tracer,
273+
s.logContext,
268274
s.writerReconnectorID,
269275
s.sessionID,
270276
codec.ToInt32(),

0 commit comments

Comments
 (0)