fix initializing of lazy transactions + refactored traces

Author: asmyasnikov

internal/table/session.go

@@ -671,26 +671,31 @@ func (s *session) Execute(
 		return nil, nil, xerrors.WithStackTrace(err)
 	}
 
-	return s.executeQueryResult(result, txControl)
+	return s.executeQueryResult(result, request.TxControl)
 }
 
 // executeQueryResult returns Transaction and result built from received
 // result.
-func (s *session) executeQueryResult(res *Ydb_Table.ExecuteQueryResult, txControl *table.TransactionControl) (
+func (s *session) executeQueryResult(
+	res *Ydb_Table.ExecuteQueryResult, txControl *Ydb_Table.TransactionControl,
+) (
 	table.Transaction, result.Result, error,
 ) {
-	t := &transaction{
-		id:      res.GetTxMeta().GetId(),
-		state:   txStateInitialized,
-		s:       s,
-		control: txControl,
+	tx := &transaction{
+		id: res.GetTxMeta().GetId(),
+		s:  s,
 	}
-	r := scanner.NewUnary(
+	if txControl.CommitTx {
+		tx.state = txStateCommitted
+	} else {
+		tx.state = txStateInitialized
+		tx.control = table.TxControl(table.WithTxID(tx.id))
+	}
+	return tx, scanner.NewUnary(
 		res.GetResultSets(),
 		res.GetQueryStats(),
 		scanner.WithIgnoreTruncated(s.config.IgnoreTruncated()),
-	)
-	return t, r, nil
+	), nil
 }
 
 // executeDataQuery executes data query.

internal/table/statement.go

@@ -63,21 +63,20 @@ func (s *statement) Execute(
 		onDone(txr, true, r, err)
 	}()
 
-	return s.execute(ctx, a, request, txControl)
+	return s.execute(ctx, a, request, request.TxControl)
 }
 
 // execute executes prepared query without any tracing.
 func (s *statement) execute(
 	ctx context.Context, a *allocator.Allocator,
-	request *Ydb_Table.ExecuteDataQueryRequest, txControl *table.TransactionControl,
+	request *Ydb_Table.ExecuteDataQueryRequest, txControl *Ydb_Table.TransactionControl,
 ) (
 	txr table.Transaction, r result.Result, err error,
 ) {
 	res, err := s.session.executeDataQuery(ctx, a, request)
 	if err != nil {
 		return nil, nil, xerrors.WithStackTrace(err)
 	}
-
 	return s.session.executeQueryResult(res, txControl)
 }
 

internal/table/transaction.go

@@ -18,16 +18,14 @@ import (
 )
 
 var (
-	errTxInvalidatedWithCommit = xerrors.Wrap(fmt.Errorf("transaction invalidated from WithCommit() call"))
-	errTxAlreadyCommitted      = xerrors.Wrap(fmt.Errorf("transaction already committed"))
-	errTxRollbackedEarly       = xerrors.Wrap(fmt.Errorf("transaction rollbacked early"))
+	errTxAlreadyCommitted = xerrors.Wrap(fmt.Errorf("transaction already committed"))
+	errTxRollbackedEarly  = xerrors.Wrap(fmt.Errorf("transaction rollbacked early"))
 )
 
 type txState int32
 
 const (
 	txStateInitialized = iota
-	txStateInvalidatedWithCommit
 	txStateCommitted
 	txStateRollbacked
 )
@@ -49,48 +47,30 @@ func (tx *transaction) Execute(
 	query string, params *table.QueryParameters,
 	opts ...options.ExecuteDataQueryOption,
 ) (r result.Result, err error) {
+	onDone := trace.TableOnSessionTransactionExecute(
+		tx.s.config.Trace(), &ctx, tx.s, tx, queryFromText(query), params, false,
+	)
+	defer func() {
+		onDone(r, err)
+	}()
+
 	switch txState(atomic.LoadInt32((*int32)(&tx.state))) {
-	case txStateInvalidatedWithCommit:
-		return nil, xerrors.WithStackTrace(errTxInvalidatedWithCommit)
 	case txStateCommitted:
 		return nil, xerrors.WithStackTrace(errTxAlreadyCommitted)
 	case txStateRollbacked:
 		return nil, xerrors.WithStackTrace(errTxRollbackedEarly)
 	default:
-	}
-	var (
-		a          = allocator.New()
-		q          = queryFromText(query)
-		optsResult options.ExecuteDataQueryDesc
-	)
-	defer a.Free()
+		_, r, err = tx.s.Execute(ctx, tx.control, query, params, opts...)
+		if err != nil {
+			return nil, xerrors.WithStackTrace(err)
+		}
 
-	if params == nil {
-		params = table.NewQueryParameters()
-	}
+		if tx.control.Desc().CommitTx {
+			atomic.StoreInt32((*int32)(&tx.state), txStateCommitted)
+		}
 
-	for _, f := range opts {
-		f(&optsResult, a)
+		return r, nil
 	}
-
-	onDone := trace.TableOnSessionTransactionExecute(
-		tx.s.config.Trace(),
-		&ctx,
-		tx.s,
-		tx,
-		q,
-		params,
-		optsResult.QueryCachePolicy.GetKeepInCache(),
-	)
-	defer func() {
-		onDone(r, err)
-	}()
-	_, r, err = tx.s.Execute(ctx, tx.control, query, params, opts...)
-	if err != nil {
-		return nil, xerrors.WithStackTrace(err)
-	}
-
-	return r, nil
 }
 
 // ExecuteStatement executes prepared statement stmt within transaction tx.
@@ -99,15 +79,6 @@ func (tx *transaction) ExecuteStatement(
 	stmt table.Statement, params *table.QueryParameters,
 	opts ...options.ExecuteDataQueryOption,
 ) (r result.Result, err error) {
-	switch txState(atomic.LoadInt32((*int32)(&tx.state))) {
-	case txStateInvalidatedWithCommit:
-		return nil, xerrors.WithStackTrace(errTxInvalidatedWithCommit)
-	case txStateCommitted:
-		return nil, xerrors.WithStackTrace(errTxAlreadyCommitted)
-	case txStateRollbacked:
-		return nil, xerrors.WithStackTrace(errTxRollbackedEarly)
-	default:
-	}
 	if params == nil {
 		params = table.NewQueryParameters()
 	}
@@ -126,29 +97,29 @@ func (tx *transaction) ExecuteStatement(
 		&ctx,
 		tx.s,
 		tx,
+		stmt.(*statement).query,
 		params,
 	)
 	defer func() {
 		onDone(r, err)
 	}()
 
-	_, r, err = stmt.Execute(ctx, tx.control, params, opts...)
-	if err != nil {
-		return nil, xerrors.WithStackTrace(err)
-	}
+	switch txState(atomic.LoadInt32((*int32)(&tx.state))) {
+	case txStateCommitted:
+		return nil, xerrors.WithStackTrace(errTxAlreadyCommitted)
+	case txStateRollbacked:
+		return nil, xerrors.WithStackTrace(errTxRollbackedEarly)
+	default:
+		_, r, err = stmt.Execute(ctx, tx.control, params, opts...)
+		if err != nil {
+			return nil, xerrors.WithStackTrace(err)
+		}
 
-	return r, nil
-}
+		if tx.control.Desc().CommitTx {
+			atomic.StoreInt32((*int32)(&tx.state), txStateCommitted)
+		}
 
-func (tx *transaction) WithCommit() table.TransactionActor {
-	defer func() {
-		atomic.StoreInt32((*int32)(&tx.state), txStateInvalidatedWithCommit)
-	}()
-	return &transaction{
-		id:      tx.id,
-		s:       tx.s,
-		control: table.TxControl(table.WithTxID(tx.id), table.CommitTx()),
-		state:   txState(atomic.LoadInt32((*int32)(&tx.state))),
+		return r, nil
 	}
 }
 
@@ -157,20 +128,6 @@ func (tx *transaction) CommitTx(
 	ctx context.Context,
 	opts ...options.CommitTransactionOption,
 ) (r result.Result, err error) {
-	switch txState(atomic.LoadInt32((*int32)(&tx.state))) {
-	case txStateInvalidatedWithCommit:
-		return nil, xerrors.WithStackTrace(errTxInvalidatedWithCommit)
-	case txStateCommitted:
-		return nil, xerrors.WithStackTrace(errTxAlreadyCommitted)
-	case txStateRollbacked:
-		return nil, xerrors.WithStackTrace(errTxRollbackedEarly)
-	default:
-		defer func() {
-			if err == nil {
-				atomic.StoreInt32((*int32)(&tx.state), txStateCommitted)
-			}
-		}()
-	}
 	onDone := trace.TableOnSessionTransactionCommit(
 		tx.s.config.Trace(),
 		&ctx,
@@ -180,56 +137,54 @@ func (tx *transaction) CommitTx(
 	defer func() {
 		onDone(err)
 	}()
-	var (
-		request = &Ydb_Table.CommitTransactionRequest{
-			SessionId: tx.s.id,
-			TxId:      tx.id,
-			OperationParams: operation.Params(
-				ctx,
-				tx.s.config.OperationTimeout(),
-				tx.s.config.OperationCancelAfter(),
-				operation.ModeSync,
-			),
+
+	switch txState(atomic.LoadInt32((*int32)(&tx.state))) {
+	case txStateCommitted:
+		return nil, xerrors.WithStackTrace(errTxAlreadyCommitted)
+	case txStateRollbacked:
+		return nil, xerrors.WithStackTrace(errTxRollbackedEarly)
+	default:
+		var (
+			request = &Ydb_Table.CommitTransactionRequest{
+				SessionId: tx.s.id,
+				TxId:      tx.id,
+				OperationParams: operation.Params(
+					ctx,
+					tx.s.config.OperationTimeout(),
+					tx.s.config.OperationCancelAfter(),
+					operation.ModeSync,
+				),
+			}
+			response *Ydb_Table.CommitTransactionResponse
+			result   = new(Ydb_Table.CommitTransactionResult)
+		)
+
+		for _, opt := range opts {
+			opt((*options.CommitTransactionDesc)(request))
 		}
-		response *Ydb_Table.CommitTransactionResponse
-		result   = new(Ydb_Table.CommitTransactionResult)
-	)
 
-	for _, opt := range opts {
-		opt((*options.CommitTransactionDesc)(request))
-	}
+		response, err = tx.s.tableService.CommitTransaction(ctx, request)
+		if err != nil {
+			return nil, xerrors.WithStackTrace(err)
+		}
 
-	response, err = tx.s.tableService.CommitTransaction(ctx, request)
-	if err != nil {
-		return nil, xerrors.WithStackTrace(err)
-	}
-	err = response.GetOperation().GetResult().UnmarshalTo(result)
-	if err != nil {
-		return nil, xerrors.WithStackTrace(err)
+		err = response.GetOperation().GetResult().UnmarshalTo(result)
+		if err != nil {
+			return nil, xerrors.WithStackTrace(err)
+		}
+
+		atomic.StoreInt32((*int32)(&tx.state), txStateCommitted)
+
+		return scanner.NewUnary(
+			nil,
+			result.GetQueryStats(),
+			scanner.WithIgnoreTruncated(tx.s.config.IgnoreTruncated()),
+		), nil
 	}
-	return scanner.NewUnary(
-		nil,
-		result.GetQueryStats(),
-		scanner.WithIgnoreTruncated(tx.s.config.IgnoreTruncated()),
-	), nil
 }
 
 // Rollback performs a rollback of the specified active transaction.
 func (tx *transaction) Rollback(ctx context.Context) (err error) {
-	switch txState(atomic.LoadInt32((*int32)(&tx.state))) {
-	case txStateInvalidatedWithCommit:
-		return xerrors.WithStackTrace(errTxInvalidatedWithCommit)
-	case txStateCommitted:
-		return xerrors.WithStackTrace(errTxAlreadyCommitted)
-	case txStateRollbacked:
-		return xerrors.WithStackTrace(errTxRollbackedEarly)
-	default:
-		defer func() {
-			if err == nil {
-				atomic.StoreInt32((*int32)(&tx.state), txStateRollbacked)
-			}
-		}()
-	}
 	onDone := trace.TableOnSessionTransactionRollback(
 		tx.s.config.Trace(),
 		&ctx,
@@ -240,17 +195,30 @@ func (tx *transaction) Rollback(ctx context.Context) (err error) {
 		onDone(err)
 	}()
 
-	_, err = tx.s.tableService.RollbackTransaction(ctx,
-		&Ydb_Table.RollbackTransactionRequest{
-			SessionId: tx.s.id,
-			TxId:      tx.id,
-			OperationParams: operation.Params(
-				ctx,
-				tx.s.config.OperationTimeout(),
-				tx.s.config.OperationCancelAfter(),
-				operation.ModeSync,
-			),
-		},
-	)
-	return xerrors.WithStackTrace(err)
+	switch txState(atomic.LoadInt32((*int32)(&tx.state))) {
+	case txStateCommitted:
+		return nil // nop for committed tx
+	case txStateRollbacked:
+		return xerrors.WithStackTrace(errTxRollbackedEarly)
+	default:
+		_, err = tx.s.tableService.RollbackTransaction(ctx,
+			&Ydb_Table.RollbackTransactionRequest{
+				SessionId: tx.s.id,
+				TxId:      tx.id,
+				OperationParams: operation.Params(
+					ctx,
+					tx.s.config.OperationTimeout(),
+					tx.s.config.OperationCancelAfter(),
+					operation.ModeSync,
+				),
+			},
+		)
+		if err != nil {
+			return xerrors.WithStackTrace(err)
+		}
+
+		atomic.StoreInt32((*int32)(&tx.state), txStateRollbacked)
+
+		return nil
+	}
 }

internal/table/transaction_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Table"
 	"google.golang.org/protobuf/proto"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
 	"github.com/ydb-platform/ydb-go-sdk/v3/testutil"
 )
@@ -113,11 +114,11 @@ func TestTxSkipRollbackForCommitted(t *testing.T) {
 			t.Fatalf("unexpected rollback: %d", begin)
 		}
 		_, err = x.CommitTx(context.Background())
-		if err != nil {
-			t.Fatal(err)
+		if !xerrors.Is(err, errTxRollbackedEarly) {
+			t.Fatal("must be errTxRollbackedEarly")
 		}
-		if commit != 2 {
-			t.Fatalf("unexpected commit: %d", begin)
+		if commit != 1 {
+			t.Fatalf("unexpected commit: %d", commit)
 		}
 	}
 }