github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/topic/topicwriterinternal/writer_transaction.go (about)

     1  package topicwriterinternal
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
     7  	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
     8  )
     9  
    10  type WriterWithTransaction struct {
    11  	streamWriter *WriterReconnector
    12  	tx           tx.Transaction
    13  	tracer       *trace.Topic
    14  }
    15  
    16  func NewTopicWriterTransaction(w *WriterReconnector, tx tx.Transaction, tracer *trace.Topic) *WriterWithTransaction {
    17  	res := &WriterWithTransaction{
    18  		streamWriter: w,
    19  		tx:           tx,
    20  	}
    21  	if tracer == nil {
    22  		res.tracer = &trace.Topic{}
    23  	} else {
    24  		res.tracer = tracer
    25  	}
    26  
    27  	tx.OnBeforeCommit(res.onBeforeCommitTransaction)
    28  	tx.OnCompleted(res.onTransactionCompleted)
    29  
    30  	return res
    31  }
    32  
    33  func (w *WriterWithTransaction) onBeforeCommitTransaction(ctx context.Context) (err error) {
    34  	traceCtx := ctx
    35  	onDone := trace.TopicOnWriterBeforeCommitTransaction(
    36  		w.tracer,
    37  		&traceCtx,
    38  		w.tx.SessionID(),
    39  		w.streamWriter.GetSessionID(),
    40  		w.tx.ID(),
    41  	)
    42  	ctx = traceCtx
    43  
    44  	defer func() {
    45  		onDone(err, w.streamWriter.GetSessionID())
    46  	}()
    47  
    48  	// wait message flushing
    49  	return w.Close(ctx)
    50  }
    51  
    52  func (w *WriterWithTransaction) WaitInit(ctx context.Context) (info InitialInfo, err error) {
    53  	return w.streamWriter.WaitInit(ctx)
    54  }
    55  
    56  func (w *WriterWithTransaction) Write(ctx context.Context, messages ...PublicMessage) error {
    57  	for i := range messages {
    58  		messages[i].tx = w.tx
    59  	}
    60  
    61  	return w.streamWriter.Write(ctx, messages)
    62  }
    63  
    64  func (w *WriterWithTransaction) Close(ctx context.Context) error {
    65  	return w.streamWriter.Close(ctx)
    66  }
    67  
    68  func (w *WriterWithTransaction) onTransactionCompleted(err error) {
    69  	// after transaction finished by any reason - the writer closed without flush
    70  	noNeedFlushCtx, cancel := context.WithCancel(context.Background())
    71  	cancel()
    72  
    73  	_ = w.Close(noNeedFlushCtx)
    74  }