fixed topic.UpdateOffsetsInTransaction (#1864)

Author: kprokopenko

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed an issue where `topic.UpdateOffsetsInTransaction` was executed on a node different from the one where the transaction was running, which could lead to the error "Database coordinators are unavailable"
+
 ## v3.116.0
 * Added experimental support for query results in `Apache Arrow`
 * Fix flaky unit test TestUnboundedChanContextTimeout

internal/query/transaction.go

@@ -179,6 +179,10 @@ func (tx *Transaction) SessionID() string {
 	return tx.s.ID()
 }
 
+func (tx *Transaction) NodeID() uint32 {
+	return tx.s.NodeID()
+}
+
 func (tx *Transaction) txControl() *baseTx.Control {
 	if tx.ID() != baseTx.LazyTxID {
 		return baseTx.NewControl(baseTx.WithTxID(tx.ID()))

internal/topic/topicreaderinternal/stream_reader_impl.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/credentials"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/background"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopiccommon"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopicreader"
@@ -245,6 +246,10 @@ func (r *topicStreamReaderImpl) commitWithTransaction(
 			onDone(err)
 		}()
 
+		// UpdateOffsetsInTransaction operation must be executed on the same Node where the transaction was initiated.
+		// Otherwise, errors such as `Database coordinators are unavailable` may occur.
+		ctx = endpoint.WithNodeID(ctx, tx.NodeID())
+
 		err = r.topicClient.UpdateOffsetsInTransaction(ctx, req)
 
 		return err

internal/topic/topicreaderinternal/stream_reader_impl_test.go

@@ -14,6 +14,7 @@ import (
 	"go.uber.org/mock/gomock"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopiccommon"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopicreader"
@@ -1420,4 +1421,27 @@ func TestUpdateCommitInTransaction(t *testing.T) {
 		require.True(t, txMock.RolledBack)
 		require.True(t, txMock.materialized)
 	})
+	t.Run("UpdateOffsetsInTransaction must be executed on the tx Node", func(t *testing.T) {
+		e := newTopicReaderTestEnv(t)
+		e.Start()
+
+		txMock := newMockTransactionWrapper("test-session-id", "test-tx-id")
+		txMock.nodeID = 123
+
+		e.TopicClient.EXPECT().UpdateOffsetsInTransaction(gomock.Any(), gomock.Any()).DoAndReturn(
+			func(ctx context.Context, _ *rawtopic.UpdateOffsetsInTransactionRequest) error {
+				nodeID, ok := endpoint.ContextNodeID(ctx)
+
+				require.True(t, ok)
+				require.Equal(t, uint32(123), nodeID)
+
+				return nil
+			})
+
+		batch, err := topicreadercommon.NewBatch(e.partitionSession, nil)
+		require.NoError(t, err)
+
+		err = e.reader.commitWithTransaction(e.ctx, txMock, batch)
+		require.NoError(t, err)
+	})
 }

internal/topic/topicreaderinternal/transaction_mock_test.go

@@ -21,6 +21,7 @@ type mockTransaction struct {
 	materializedID tx.Identifier
 	materialized   bool
 	sessionID      string
+	nodeID         uint32
 	onCompleted    []tx.OnTransactionCompletedFunc
 	RolledBack     bool
 }
@@ -36,6 +37,10 @@ func (m *mockTransaction) SessionID() string {
 	return m.sessionID
 }
 
+func (m *mockTransaction) NodeID() uint32 {
+	return m.nodeID
+}
+
 func (m *mockTransaction) OnBeforeCommit(f tx.OnTransactionBeforeCommit) {
 }
 

internal/tx/transaction.go

@@ -11,6 +11,7 @@ type Transaction interface {
 	Identifier
 	UnLazy(ctx context.Context) error
 	SessionID() string
+	NodeID() uint32
 
 	// OnBeforeCommit add callback, which will be called before commit transaction
 	// the method will be not call the method if some error happen and transaction will not be committed