test: added more transaction resolver test

Signed-off-by: Harshit Gangal <harshit@planetscale.com>

go/test/endtoend/transaction/twopc/main_test.go

@@ -213,6 +213,10 @@ func runVStream(t *testing.T, ctx context.Context, ch chan *binlogdatapb.VEvent,
 }
 
 func retrieveTransitions(t *testing.T, ch chan *binlogdatapb.VEvent, tableMap map[string][]*querypb.Field, dtMap map[string]string) map[string][]string {
+	return retrieveTransitionsWithTimeout(t, ch, tableMap, dtMap, 1*time.Second)
+}
+
+func retrieveTransitionsWithTimeout(t *testing.T, ch chan *binlogdatapb.VEvent, tableMap map[string][]*querypb.Field, dtMap map[string]string, timeout time.Duration) map[string][]string {
 	logTable := make(map[string][]string)
 
 	keepWaiting := true
@@ -231,7 +235,7 @@ func retrieveTransitions(t *testing.T, ch chan *binlogdatapb.VEvent, tableMap ma
 			if re.FieldEvent != nil {
 				tableMap[re.FieldEvent.TableName] = re.FieldEvent.Fields
 			}
-		case <-time.After(2 * time.Second):
+		case <-time.After(timeout):
 			keepWaiting = false
 		}
 	}

go/test/endtoend/transaction/twopc/twopc_test.go

@@ -23,6 +23,7 @@ import (
 	"sort"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -553,16 +554,15 @@ func TestDTResolveAfterMMCommit(t *testing.T) {
 	_, err = conn.Execute(qCtx, "insert into twopc_user(id, name) values(10,'apa')", nil)
 	require.NoError(t, err)
 
-	// This setting will return from VTGate after commit happens on metadata manager.
-	// The RMs will be left in PREPARE state to resolve.
+	// The caller ID is used to simulate the failure at the desired point.
 	newCtx := callerid.NewContext(qCtx, callerid.NewEffectiveCallerID("MMCommitted_FailNow", "", ""), nil)
 	_, err = conn.Execute(newCtx, "commit", nil)
-	require.NoError(t, err)
+	require.ErrorContains(t, err, "Fail After MM commit")
 
 	// Below check ensures that the transaction is resolved by the resolver on receiving unresolved transaction signal from MM.
 	tableMap := make(map[string][]*querypb.Field)
 	dtMap := make(map[string]string)
-	logTable := retrieveTransitions(t, ch, tableMap, dtMap)
+	logTable := retrieveTransitionsWithTimeout(t, ch, tableMap, dtMap, 2*time.Second)
 	expectations := map[string][]string{
 		"ks.dt_state:80-": {
 			"insert:[VARCHAR(\"dtid-1\") VARCHAR(\"PREPARE\")]",
@@ -599,8 +599,6 @@ func TestDTResolveAfterMMCommit(t *testing.T) {
 // TestDTResolveAfterPrepare tests that transaction is rolled back on recovery
 // failure after RM prepare and before MM commit.
 func TestDTResolveAfterPrepare(t *testing.T) {
-	defer cleanup(t)
-
 	vtgateConn, err := cluster.DialVTGate(context.Background(), t.Name(), vtgateGrpcAddress, "dt_user", "")
 	require.NoError(t, err)
 	defer vtgateConn.Close()
@@ -622,16 +620,15 @@ func TestDTResolveAfterPrepare(t *testing.T) {
 	_, err = conn.Execute(qCtx, "insert into twopc_user(id, name) values(8,'bar')", nil)
 	require.NoError(t, err)
 
-	// This setting will return from VTGate after commit happens on metadata manager.
-	// The RMs will be left in PREPARE state to resolve.
+	// The caller ID is used to simulate the failure at the desired point.
 	newCtx := callerid.NewContext(qCtx, callerid.NewEffectiveCallerID("RMPrepared_FailNow", "", ""), nil)
 	_, err = conn.Execute(newCtx, "commit", nil)
 	require.ErrorContains(t, err, "Fail After RM prepared")
 
 	// Below check ensures that the transaction is resolved by the resolver on receiving unresolved transaction signal from MM.
 	tableMap := make(map[string][]*querypb.Field)
 	dtMap := make(map[string]string)
-	logTable := retrieveTransitions(t, ch, tableMap, dtMap)
+	logTable := retrieveTransitionsWithTimeout(t, ch, tableMap, dtMap, 2*time.Second)
 	expectations := map[string][]string{
 		"ks.dt_state:80-": {
 			"insert:[VARCHAR(\"dtid-1\") VARCHAR(\"PREPARE\")]",
@@ -654,3 +651,51 @@ func TestDTResolveAfterPrepare(t *testing.T) {
 	assert.Equal(t, expectations, logTable,
 		"mismatch expected: \n got: %s, want: %s", prettyPrint(logTable), prettyPrint(expectations))
 }
+
+// TestDTResolveAfterTransactionRecord tests that transaction is rolled back on recovery
+// failure after TR created and before RM prepare.
+func TestDTResolveAfterTransactionRecord(t *testing.T) {
+	vtgateConn, err := cluster.DialVTGate(context.Background(), t.Name(), vtgateGrpcAddress, "dt_user", "")
+	require.NoError(t, err)
+	defer vtgateConn.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	ch := make(chan *binlogdatapb.VEvent)
+	runVStream(t, ctx, ch, vtgateConn)
+
+	conn := vtgateConn.Session("", nil)
+	qCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	// Insert into multiple shards
+	_, err = conn.Execute(qCtx, "begin", nil)
+	require.NoError(t, err)
+	_, err = conn.Execute(qCtx, "insert into twopc_user(id, name) values(7,'foo')", nil)
+	require.NoError(t, err)
+	_, err = conn.Execute(qCtx, "insert into twopc_user(id, name) values(8,'bar')", nil)
+	require.NoError(t, err)
+
+	// The caller ID is used to simulate the failure at the desired point.
+	newCtx := callerid.NewContext(qCtx, callerid.NewEffectiveCallerID("TRCreated_FailNow", "", ""), nil)
+	_, err = conn.Execute(newCtx, "commit", nil)
+	require.ErrorContains(t, err, "Fail After TR created")
+
+	// Below check ensures that the transaction is resolved by the resolver on receiving unresolved transaction signal from MM.
+	tableMap := make(map[string][]*querypb.Field)
+	dtMap := make(map[string]string)
+	logTable := retrieveTransitionsWithTimeout(t, ch, tableMap, dtMap, 2*time.Second)
+	expectations := map[string][]string{
+		"ks.dt_state:80-": {
+			"insert:[VARCHAR(\"dtid-1\") VARCHAR(\"PREPARE\")]",
+			"update:[VARCHAR(\"dtid-1\") VARCHAR(\"ROLLBACK\")]",
+			"delete:[VARCHAR(\"dtid-1\") VARCHAR(\"ROLLBACK\")]",
+		},
+		"ks.dt_participant:80-": {
+			"insert:[VARCHAR(\"dtid-1\") INT64(1) VARCHAR(\"ks\") VARCHAR(\"-80\")]",
+			"delete:[VARCHAR(\"dtid-1\") INT64(1) VARCHAR(\"ks\") VARCHAR(\"-80\")]",
+		},
+	}
+	assert.Equal(t, expectations, logTable,
+		"mismatch expected: \n got: %s, want: %s", prettyPrint(logTable), prettyPrint(expectations))
+}

go/vt/vtgate/tx_conn.go

@@ -201,6 +201,14 @@ func (txc *TxConn) commit2PC(ctx context.Context, session *SafeSession) error {
 		return err
 	}
 
+	// retrieve the caller ID to execute the testing flow
+	callerID := callerid.EffectiveCallerIDFromContext(ctx)
+
+	// Test code to simulate a failure after RM prepare
+	if failNow, err := checkTestFailure(callerID, "TRCreated_FailNow"); failNow {
+		return err
+	}
+
 	err = txc.runSessions(ctx, session.ShardSessions[1:], session.logging, func(ctx context.Context, s *vtgatepb.Session_ShardSession, logging *executeLogger) error {
 		return txc.tabletGateway.Prepare(ctx, s.Target, s.TransactionId, dtid)
 	})
@@ -214,9 +222,6 @@ func (txc *TxConn) commit2PC(ctx context.Context, session *SafeSession) error {
 		return err
 	}
 
-	// retrieve the caller ID to execute the testing flow
-	callerID := callerid.EffectiveCallerIDFromContext(ctx)
-
 	// Test code to simulate a failure after RM prepare
 	if failNow, err := checkTestFailure(callerID, "RMPrepared_FailNow"); failNow {
 		return err
@@ -247,14 +252,18 @@ func checkTestFailure(callerID *vtrpcpb.CallerID, expectCaller string) (bool, er
 		return false, nil
 	}
 	switch callerID.Principal {
+	case "TRCreated_FailNow":
+		log.Errorf("Fail After TR created")
+		// no commit decision is made. Transaction should be a rolled back.
+		return true, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "Fail After TR created")
 	case "RMPrepared_FailNow":
 		log.Errorf("Fail After RM prepared")
 		// no commit decision is made. Transaction should be a rolled back.
 		return true, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "Fail After RM prepared")
 	case "MMCommitted_FailNow":
 		log.Errorf("Fail After MM commit")
 		//  commit decision is made. Transaction should be committed.
-		return true, nil
+		return true, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "Fail After MM commit")
 	default:
 		return false, nil
 	}