logs and traces for write messages in transaction

Author: rekby

examples/topic/topicwriter/topic_writer_transaction.go

@@ -56,7 +56,7 @@ func CopyMessagesBetweenTopics(ctx context.Context, db *ydb.Driver, reader *topi
 
 func CopyMessagesBetweenTopicsTxWriter(ctx context.Context, db *ydb.Driver, reader *topicreader.Reader, topic string) error {
 	return db.Query().DoTx(ctx, func(ctx context.Context, tx query.TxActor) error {
-		writer, err := db.Topic().StartWransactionalWriter(ctx, tx, topic)
+		writer, err := db.Topic().StartTransactionalWriter(ctx, tx, topic)
 		if err != nil {
 			return err
 		}

internal/grpcwrapper/rawtopic/rawtopicwriter/messages.go

@@ -166,10 +166,6 @@ func (r *WriteRequest) toProto() (p *Ydb_Topic.StreamWriteMessage_FromClient_Wri
 		},
 	}
 
-	if r.Tx.ID != "" || r.Tx.Session != "" {
-		res.WriteRequest.Tx = &Ydb_Topic.TransactionIdentity{}
-	}
-
 	return res, nil
 }
 

internal/topic/topicclientinternal/client.go

@@ -2,6 +2,7 @@ package topicclientinternal
 
 import (
 	"context"
+	"errors"
 
 	"github.com/ydb-platform/ydb-go-genproto/Ydb_Topic_V1"
 	"google.golang.org/grpc"
@@ -14,6 +15,8 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topicreadercommon"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topicreaderinternal"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topicwriterinternal"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
 	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topiclistener"
 	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topicoptions"
@@ -23,6 +26,10 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
+var (
+	errUnsupportedTransactionType = xerrors.Wrap(errors.New("ydb: unsuppotred transaction type. Use transaction from Driver().Query().DoTx(...)"))
+)
+
 type Client struct {
 	cfg                    topic.Config
 	cred                   credentials.Credentials
@@ -269,6 +276,32 @@ func (c *Client) StartReader(
 
 // StartWriter create new topic writer wrapper
 func (c *Client) StartWriter(topicPath string, opts ...topicoptions.WriterOption) (*topicwriter.Writer, error) {
+	cfg := c.createWriterConfig(topicPath, opts)
+	writer, err := topicwriterinternal.NewWriterReconnector(cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	return topicwriter.NewWriter(writer), nil
+}
+
+func (c *Client) StartTransactionalWriter(transaction tx.Identifier, topicpath string, opts ...topicoptions.WriterOption) (*topicwriter.TxWriter, error) {
+	internalTx, ok := transaction.(tx.Transaction)
+	if !ok {
+		return nil, xerrors.WithStackTrace(errUnsupportedTransactionType)
+	}
+
+	cfg := c.createWriterConfig(topicpath, opts)
+	writer, err := topicwriterinternal.NewWriterReconnector(cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	txWriter := topicwriterinternal.NewTopicWriterTransaction(writer, internalTx, cfg.Tracer)
+	return topicwriter.NewTxWriterInternal(txWriter), nil
+}
+
+func (c *Client) createWriterConfig(topicPath string, opts []topicoptions.WriterOption) topicwriterinternal.WriterReconnectorConfig {
 	var connector topicwriterinternal.ConnectFunc = func(ctx context.Context) (
 		topicwriterinternal.RawTopicWriterStream,
 		error,
@@ -281,14 +314,9 @@ func (c *Client) StartWriter(topicPath string, opts ...topicoptions.WriterOption
 		topicwriterinternal.WithTopic(topicPath),
 		topicwriterinternal.WithCommonConfig(c.cfg.Common),
 		topicwriterinternal.WithTrace(c.cfg.Trace),
+		topicwriterinternal.WithCredentials(c.cred),
 	}
 
 	options = append(options, opts...)
-
-	writer, err := topicwriterinternal.NewWriter(c.cred, options)
-	if err != nil {
-		return nil, err
-	}
-
-	return topicwriter.NewWriter(writer), nil
+	return topicwriterinternal.NewWriterReconnectorConfig(options...)
 }

internal/topic/topicreaderinternal/stream_reader_impl.go

@@ -297,8 +297,6 @@ func (r *topicStreamReaderImpl) addOnTransactionCompletedHandler(
 		defer onDone()
 
 		ctx = traceCtx
-		//nolint:godox
-		// TODO: trace
 		if transactionResult == nil {
 			topicreadercommon.BatchGetPartitionSession(batch).SetCommittedOffsetForward(commitRange.CommitOffsetEnd)
 		} else {

internal/topic/topicwriterinternal/message.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopiccommon"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopicwriter"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 )
 
@@ -20,6 +21,8 @@ type PublicMessage struct {
 	Data      io.Reader
 	Metadata  map[string][]byte
 
+	tx tx.Transaction
+
 	// partitioning at level message available by protocol, but doesn't available by current server implementation
 	// the field hidden from public access for prevent runtime errors.
 	// it will be published after implementation on server side.

internal/topic/topicwriterinternal/writer.go

@@ -18,7 +18,7 @@ type WritersCommonConfig struct {
 	defaultPartitioning rawtopicwriter.Partitioning
 	compressorCount     int
 
-	tracer             *trace.Topic
+	Tracer             *trace.Topic
 	cred               credentials.Credentials
 	credUpdateInterval time.Duration
 	clock              clockwork.Clock

internal/topic/topicwriterinternal/writer_config.go

@@ -145,7 +145,7 @@ func WithWaitAckOnWrite(val bool) PublicWriterOption {
 
 func WithTrace(tracer *trace.Topic) PublicWriterOption {
 	return func(cfg *WriterReconnectorConfig) {
-		cfg.tracer = cfg.tracer.Compose(tracer)
+		cfg.Tracer = cfg.Tracer.Compose(tracer)
 	}
 }
 

internal/topic/topicwriterinternal/writer_options.go

@@ -38,6 +38,7 @@ var (
 	errNoAllowedCodecs       = xerrors.Wrap(errors.New("ydb: no allowed codecs for write to topic"))
 	errLargeMessage          = xerrors.Wrap(errors.New("ydb: message uncompressed size more, then limit"))
 	PublicErrQueueIsFull     = xerrors.Wrap(errors.New("ydb: queue is full"))
+	errDiffetentTransactions = xerrors.Wrap(errors.New("ydb: internal writer has messages from different trasactions. It is internal logic error, write issue please: https://github.com/ydb-platform/ydb-go-sdk/issues/new?assignees=&labels=bug&projects=&template=01_BUG_REPORT.md&title=bug%3A+"))
 
 	// errProducerIDNotEqualMessageGroupID is temporary
 	// WithMessageGroupID is optional parameter because it allowed to be skipped by protocol.
@@ -73,14 +74,14 @@ func (cfg *WriterReconnectorConfig) validate() error {
 	return nil
 }
 
-func newWriterReconnectorConfig(options ...PublicWriterOption) WriterReconnectorConfig {
+func NewWriterReconnectorConfig(options ...PublicWriterOption) WriterReconnectorConfig {
 	cfg := WriterReconnectorConfig{
 		WritersCommonConfig: WritersCommonConfig{
 			cred:               credentials.NewAnonymousCredentials(),
 			credUpdateInterval: time.Hour,
 			clock:              clockwork.NewRealClock(),
 			compressorCount:    runtime.NumCPU(),
-			tracer:             &trace.Topic{},
+			Tracer:             &trace.Topic{},
 		},
 		AutoSetSeqNo:       true,
 		AutoSetCreatedTime: true,
@@ -119,25 +120,29 @@ type WriterReconnector struct {
 	background                     background.Worker
 	retrySettings                  topic.RetrySettings
 	writerInstanceID               string
-	sessionID                      string
 	semaphore                      *semaphore.Weighted
 	firstInitResponseProcessedChan empty.Chan
 	lastSeqNo                      int64
 	encodersMap                    *EncoderMap
 	initDoneCh                     empty.Chan
 	initInfo                       InitialInfo
 	m                              xsync.RWMutex
+	sessionID                      string
 	firstConnectionHandled         atomic.Bool
 	initDone                       bool
 }
 
-func newWriterReconnector(
-	cfg WriterReconnectorConfig, //nolint:gocritic
-) *WriterReconnector {
+func NewWriterReconnector(
+	cfg WriterReconnectorConfig,
+) (*WriterReconnector, error) {
+	if err := cfg.validate(); err != nil {
+		return nil, err
+	}
+
 	res := newWriterReconnectorStopped(cfg)
 	res.start()
 
-	return res
+	return res, nil
 }
 
 func newWriterReconnectorStopped(
@@ -313,7 +318,7 @@ func (w *WriterReconnector) createMessagesWithContent(messages []PublicMessage)
 		sessionID = w.sessionID
 	})
 	onCompressDone := trace.TopicOnWriterCompressMessages(
-		w.cfg.tracer,
+		w.cfg.Tracer,
 		w.writerInstanceID,
 		sessionID,
 		w.cfg.forceCodec.ToInt32(),
@@ -354,7 +359,7 @@ func (w *WriterReconnector) Close(ctx context.Context) error {
 }
 
 func (w *WriterReconnector) close(ctx context.Context, reason error) (resErr error) {
-	onDone := trace.TopicOnWriterClose(w.cfg.tracer, w.writerInstanceID, reason)
+	onDone := trace.TopicOnWriterClose(w.cfg.Tracer, w.writerInstanceID, reason)
 	defer func() {
 		onDone(resErr)
 	}()
@@ -461,7 +466,7 @@ func (w *WriterReconnector) startWriteStream(ctx, streamCtx context.Context, att
 	err error,
 ) {
 	traceOnDone := trace.TopicOnWriterReconnect(
-		w.cfg.tracer,
+		w.cfg.Tracer,
 		w.writerInstanceID,
 		w.cfg.topic,
 		w.cfg.producerID,
@@ -625,6 +630,13 @@ func (w *WriterReconnector) createWriterStreamConfig(stream RawTopicWriterStream
 	return cfg
 }
 
+func (w *WriterReconnector) GetSessionID() (sessionID string) {
+	w.m.WithLock(func() {
+		sessionID = w.sessionID
+	})
+	return sessionID
+}
+
 func sendMessagesToStream(
 	stream RawTopicWriterStream,
 	targetCodec rawtopiccommon.Codec,
@@ -684,6 +696,17 @@ func createWriteRequest(messages []messageWithDataContent, targetCodec rawtopicc
 	res rawtopicwriter.WriteRequest,
 	err error,
 ) {
+	for i := 1; i < len(messages); i++ {
+		if messages[i-1].tx != messages[i].tx {
+			return res, xerrors.WithStackTrace(errDiffetentTransactions)
+		}
+	}
+
+	if len(messages) > 0 {
+		res.Tx.ID = messages[0].tx.ID()
+		res.Tx.Session = messages[0].tx.SessionID()
+	}
+
 	res.Codec = targetCodec
 	res.Messages = make([]rawtopicwriter.MessageData, len(messages))
 	for i := range messages {

internal/topic/topicwriterinternal/writer_reconnector.go

@@ -31,7 +31,7 @@ var testCommonEncoders = NewEncoderMap()
 func TestWriterImpl_AutoSeq(t *testing.T) {
 	t.Run("OK", func(t *testing.T) {
 		ctx := xtest.Context(t)
-		w := newWriterReconnectorStopped(newWriterReconnectorConfig(
+		w := newWriterReconnectorStopped(NewWriterReconnectorConfig(
 			WithAutoSetSeqNo(true),
 			WithAutosetCreatedTime(false),
 		))
@@ -63,7 +63,7 @@ func TestWriterImpl_AutoSeq(t *testing.T) {
 	t.Run("PredefinedSeqNo", func(t *testing.T) {
 		ctx := xtest.Context(t)
 
-		w := newWriterReconnectorStopped(newWriterReconnectorConfig(WithAutoSetSeqNo(true)))
+		w := newWriterReconnectorStopped(NewWriterReconnectorConfig(WithAutoSetSeqNo(true)))
 		w.firstConnectionHandled.Store(true)
 		require.Error(t, w.Write(ctx, newTestMessages(1)))
 	})
@@ -72,7 +72,7 @@ func TestWriterImpl_AutoSeq(t *testing.T) {
 func TestWriterImpl_CheckMessages(t *testing.T) {
 	t.Run("MessageSize", func(t *testing.T) {
 		ctx := xtest.Context(t)
-		w := newWriterReconnectorStopped(newWriterReconnectorConfig())
+		w := newWriterReconnectorStopped(NewWriterReconnectorConfig())
 		w.firstConnectionHandled.Store(true)
 
 		maxSize := 5
@@ -282,7 +282,7 @@ func TestWriterImpl_WriteCodecs(t *testing.T) {
 func TestWriterReconnector_Write_QueueLimit(t *testing.T) {
 	xtest.TestManyTimes(t, func(t testing.TB) {
 		ctx := xtest.Context(t)
-		w := newWriterReconnectorStopped(newWriterReconnectorConfig(
+		w := newWriterReconnectorStopped(NewWriterReconnectorConfig(
 			WithAutoSetSeqNo(false),
 			WithMaxQueueLen(2),
 		))
@@ -994,7 +994,7 @@ func newTestMessagesWithContent(numbers ...int) []messageWithDataContent {
 
 func newTestWriterStopped(opts ...PublicWriterOption) *WriterReconnector {
 	cfgOptions := append(defaultTestWriterOptions(), opts...)
-	cfg := newWriterReconnectorConfig(cfgOptions...)
+	cfg := NewWriterReconnectorConfig(cfgOptions...)
 	res := newWriterReconnectorStopped(cfg)
 
 	if cfg.AdditionalEncoders == nil {
@@ -1074,7 +1074,7 @@ func newTestEnv(t testing.TB, options *testEnvOptions) *testEnv {
 	}))
 	writerOptions = append(writerOptions, options.writerOptions...)
 
-	res.writer = newWriterReconnectorStopped(newWriterReconnectorConfig(writerOptions...))
+	res.writer = newWriterReconnectorStopped(NewWriterReconnectorConfig(writerOptions...))
 
 	res.stream.EXPECT().Recv().DoAndReturn(res.receiveMessageHandler).AnyTimes()