reuse txnop

pkg/frontend/mysql_protocol_test.go

@@ -186,19 +186,17 @@ func TestKill(t *testing.T) {
 	eng.EXPECT().Hints().Return(engine.Hints{CommitOrRollbackTimeout: time.Second * 10}).AnyTimes()
 
 	txnClient := mock_frontend.NewMockTxnClient(ctrl)
-	txnClient.EXPECT().New(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, commitTS any, options ...any) (client.TxnOperator, error) {
-		wp := newTestWorkspace()
-		txnOp := mock_frontend.NewMockTxnOperator(ctrl)
-		txnOp.EXPECT().Txn().Return(txn.TxnMeta{}).AnyTimes()
-		txnOp.EXPECT().GetWorkspace().Return(wp).AnyTimes()
-		txnOp.EXPECT().Commit(gomock.Any()).Return(nil).AnyTimes()
-		txnOp.EXPECT().Rollback(gomock.Any()).Return(nil).AnyTimes()
-		txnOp.EXPECT().SetFootPrints(gomock.Any(), gomock.Any()).Return().AnyTimes()
-		txnOp.EXPECT().Status().Return(txn.TxnStatus_Active).AnyTimes()
-		txnOp.EXPECT().EnterRunSql().Return().AnyTimes()
-		txnOp.EXPECT().ExitRunSql().Return().AnyTimes()
-		return txnOp, nil
-	}).AnyTimes()
+	txnClient.EXPECT().New(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
+		func(ctx context.Context, commitTS any, options ...any) (client.TxnOperator, error) {
+			txnOp := newTestTxnOp()
+			return txnOp, nil
+		}).AnyTimes()
+	txnClient.EXPECT().RestartTxn(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
+		func(ctx context.Context, txnOp TxnOperator, commitTS any, options ...any) (client.TxnOperator, error) {
+			tTxnOp := txnOp.(*testTxnOp)
+			tTxnOp.meta.Status = txn.TxnStatus_Active
+			return txnOp, nil
+		}).AnyTimes()
 	pu, err := getParameterUnit("test/system_vars_config.toml", eng, txnClient)
 	require.NoError(t, err)
 	pu.SV.SkipCheckUser = true

pkg/frontend/test/Makefile

@@ -1,10 +1,14 @@
 CURRENT_DIR = $(shell pwd)
 
 .PHONY: generate_mock
-generate_mock:
+generate_mock: ../../../pkg/txn/client/types.go ../../../pkg/vm/engine/types.go
 	@go install github.com/golang/mock/mockgen@v1.6.0
 	@echo "Current Directory " $(CURRENT_DIR)
 	@mockgen -source ../../../pkg/txn/client/types.go -package mock_frontend > txn_mock.go
 	@mockgen -source ../../../pkg/vm/engine/types.go -package mock_frontend > engine_mock.go
+
+#tricky
+generate_fe_mock: ../types.go
+	@go install github.com/golang/mock/mockgen@v1.6.0
 	@mockgen -source ../types.go -package mock_frontend > types_mock.go
 

pkg/frontend/txn.go

@@ -28,6 +28,7 @@ import (
 	moruntime "github.com/matrixorigin/matrixone/pkg/common/runtime"
 	"github.com/matrixorigin/matrixone/pkg/defines"
 	"github.com/matrixorigin/matrixone/pkg/pb/metadata"
+	"github.com/matrixorigin/matrixone/pkg/pb/txn"
 	"github.com/matrixorigin/matrixone/pkg/txn/client"
 	"github.com/matrixorigin/matrixone/pkg/txn/clock"
 	"github.com/matrixorigin/matrixone/pkg/txn/storage/memorystorage"
@@ -235,7 +236,6 @@ func (th *TxnHandler) GetTxnCtx() context.Context {
 
 // invalidateTxnUnsafe releases the txnOp and clears the server status bit SERVER_STATUS_IN_TRANS
 func (th *TxnHandler) invalidateTxnUnsafe() {
-	th.txnOp = nil
 	resetBits(&th.serverStatus, defaultServerStatus)
 	resetBits(&th.optionBits, defaultOptionBits)
 }
@@ -252,7 +252,7 @@ func (th *TxnHandler) inActiveTxnUnsafe() bool {
 	if th.txnOp != nil && th.txnCtx == nil {
 		panic("txnOp != nil and txnCtx == nil")
 	}
-	return th.txnOp != nil && th.txnCtx != nil
+	return th.txnOp != nil && th.txnOp.Txn().Status == txn.TxnStatus_Active && th.txnCtx != nil
 }
 
 // Create starts a new txn.
@@ -411,13 +411,27 @@ func (th *TxnHandler) createTxnOpUnsafe(execCtx *ExecCtx) error {
 		}
 	}
 
-	err, hasRecovered = ExecuteFuncWithRecover(func() error {
-		th.txnOp, err2 = getPu(execCtx.ses.GetService()).TxnClient.New(
-			th.txnCtx,
-			execCtx.ses.getLastCommitTS(),
-			opts...)
-		return err2
-	})
+	txnClient := getPu(execCtx.ses.GetService()).TxnClient
+	if th.txnOp == nil {
+		err, hasRecovered = ExecuteFuncWithRecover(func() error {
+			th.txnOp, err2 = txnClient.New(
+				th.txnCtx,
+				execCtx.ses.getLastCommitTS(),
+				opts...)
+			return err2
+		})
+	} else if th.txnOp.Txn().Status != txn.TxnStatus_Active {
+		err, hasRecovered = ExecuteFuncWithRecover(func() error {
+			_, err2 = txnClient.RestartTxn(
+				th.txnCtx,
+				th.txnOp,
+				execCtx.ses.getLastCommitTS(),
+				opts...)
+			return err2
+		})
+	} else {
+		return moerr.NewInternalError(execCtx.reqCtx, "NewTxnOperator: txn is already active")
+	}
 	if err != nil || hasRecovered {
 		return err
 	}