Add savepoint support to atomic distributed transaction

go/test/endtoend/cluster/cluster_process.go

@@ -170,6 +170,17 @@ func (shard *Shard) PrimaryTablet() *Vttablet {
  return shard.Vttablets[0]
 }
 
+// FindPrimaryTablet finds the primary tablet in the shard.
+func (shard *Shard) FindPrimaryTablet() *Vttablet {
+ for _, vttablet := range shard.Vttablets {
+ tabletType := vttablet.VttabletProcess.GetTabletType()
+ if tabletType == "primary" {
+ return vttablet
+ }
+ }
+ return nil
+}
+
 // Rdonly get the last tablet which is rdonly
 func (shard *Shard) Rdonly() *Vttablet {
  for idx, tablet := range shard.Vttablets {

go/test/endtoend/cluster/reshard.go

@@ -95,7 +95,7 @@ func (rw *ReshardWorkflow) WaitForVreplCatchup(timeToWait time.Duration) {
  if !slices.Contains(targetShards, shard.Name) {
  continue
  }
- vttablet := shard.PrimaryTablet().VttabletProcess
+ vttablet := shard.FindPrimaryTablet().VttabletProcess
  vttablet.WaitForVReplicationToCatchup(rw.t, rw.workflowName, fmt.Sprintf("vt_%s", vttablet.Keyspace), "", timeToWait)
  }
  }

go/test/endtoend/transaction/twopc/fuzz/fuzzer_test.go

@@ -21,6 +21,7 @@ import (
  "fmt"
  "os"
  "path"
+ "slices"
  "strconv"
  "strings"
  "sync"
@@ -126,7 +127,18 @@ func TestTwoPCFuzzTest(t *testing.T) {
  fz.start(t)
 
  // Wait for the timeForTesting so that the threads continue to run.
- time.Sleep(tt.timeForTesting)
+ timeout := time.After(tt.timeForTesting)
+ loop := true
+ for loop {
+ select {
+ case <-timeout:
+ loop = false
+ case <-time.After(1 * time.Second):
+ if t.Failed() {
+ loop = false
+ }
+ }
+ }
 
  // Signal the fuzzer to stop.
  fz.stop()
@@ -302,9 +314,11 @@ func (fz *fuzzer) generateAndExecuteTransaction(threadId int) {
  // for each update set ordered by the auto increment column will not be true.
  // That assertion depends on all the transactions running updates first to ensure that for any given update set,
  // no two transactions are running the insert queries.
- queries := []string{"begin"}
+ var queries []string
  queries = append(queries, fz.generateUpdateQueries(updateSetVal, incrementVal)...)
  queries = append(queries, fz.generateInsertQueries(updateSetVal, threadId)...)
+ queries = fz.addRandomSavePoints(queries)
+ queries = append([]string{"begin"}, queries...)
  finalCommand := "commit"
  for _, query := range queries {
  _, err := conn.ExecuteFetch(query, 0, false)
@@ -377,6 +391,45 @@ func (fz *fuzzer) runClusterDisruption(t *testing.T) {
  }
 }
 
+// addRandomSavePoints will add random savepoints and queries to the list of queries.
+// It still ensures that all the new queries added are rolledback so that the assertions of queries
+// don't change.
+func (fz *fuzzer) addRandomSavePoints(queries []string) []string {
+ savePointCount := 1
+ for {
+ shouldAddSavePoint := rand.Intn(2)
+ if shouldAddSavePoint == 0 {
+ return queries
+ }
+
+ savePointQueries := []string{"SAVEPOINT sp" + strconv.Itoa(savePointCount)}
+ randomDmlCount := rand.Intn(2) + 1
+ for i := 0; i < randomDmlCount; i++ {
+ savePointQueries = append(savePointQueries, fz.randomDML())
+ }
+ savePointQueries = append(savePointQueries, "ROLLBACK TO sp"+strconv.Itoa(savePointCount))
+ savePointCount++
+
+ savePointPosition := rand.Intn(len(queries))
+ newQueries := slices.Clone(queries[:savePointPosition])
+ newQueries = append(newQueries, savePointQueries...)
+ newQueries = append(newQueries, queries[savePointPosition:]...)
+ queries = newQueries
+ }
+}
+
+// randomDML generates a random DML to be used.
+func (fz *fuzzer) randomDML() string {
+ queryType := rand.Intn(2)
+ if queryType == 0 {
+ // Generate INSERT
+ return fmt.Sprintf(insertIntoFuzzInsert, updateRowBaseVals[rand.Intn(len(updateRowBaseVals))], rand.Intn(fz.updateSets), rand.Intn(fz.threads))
+ }
+ // Generate UPDATE
+ updateId := fz.updateRowsVals[rand.Intn(len(fz.updateRowsVals))][rand.Intn(len(updateRowBaseVals))]
+ return fmt.Sprintf(updateFuzzUpdate, rand.Intn(100000), updateId)
+}
+
 /*
 Cluster Level Disruptions for the fuzzer
 */

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

@@ -24,16 +24,19 @@ import (
  "fmt"
  "io"
  "os"
+ "strings"
  "sync"
  "testing"
  "time"
 
  "github.com/stretchr/testify/require"
 
+ "vitess.io/vitess/go/test/endtoend/utils"
+
  "vitess.io/vitess/go/mysql"
  "vitess.io/vitess/go/sqltypes"
  "vitess.io/vitess/go/test/endtoend/cluster"
- "vitess.io/vitess/go/test/endtoend/transaction/twopc/utils"
+ twopcutil "vitess.io/vitess/go/test/endtoend/transaction/twopc/utils"
  binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
  querypb "vitess.io/vitess/go/vt/proto/query"
  topodatapb "vitess.io/vitess/go/vt/proto/topodata"
@@ -42,6 +45,7 @@ import (
 
 var (
  clusterInstance *cluster.LocalProcessCluster
+ mysqlParams mysql.ConnParams
  vtParams mysql.ConnParams
  vtgateGrpcAddress string
  keyspaceName = "ks"
@@ -81,6 +85,8 @@ func TestMain(m *testing.M) {
  "--twopc_enable",
  "--twopc_abandon_age", "1",
  "--queryserver-config-transaction-cap", "3",
+ "--queryserver-config-transaction-timeout", "400s",
+ "--queryserver-config-query-timeout", "9000s",
  )
 
  // Start keyspace
@@ -102,6 +108,15 @@ func TestMain(m *testing.M) {
  vtParams = clusterInstance.GetVTParams(keyspaceName)
  vtgateGrpcAddress = fmt.Sprintf("%s:%d", clusterInstance.Hostname, clusterInstance.VtgateGrpcPort)
 
+ // create mysql instance and connection parameters
+ conn, closer, err := utils.NewMySQL(clusterInstance, keyspaceName, SchemaSQL)
+ if err != nil {
+ fmt.Println(err)
+ return 1
+ }
+ defer closer()
+ mysqlParams = conn
+
  return m.Run()
  }()
  os.Exit(exitcode)
@@ -121,8 +136,29 @@ func start(t *testing.T) (*mysql.Conn, func()) {
 
 func cleanup(t *testing.T) {
  cluster.PanicHandler(t)
- utils.ClearOutTable(t, vtParams, "twopc_user")
- utils.ClearOutTable(t, vtParams, "twopc_t1")
+ twopcutil.ClearOutTable(t, vtParams, "twopc_user")
+ twopcutil.ClearOutTable(t, vtParams, "twopc_t1")
+ sm.reset()
+}
+
+func startWithMySQL(t *testing.T) (utils.MySQLCompare, func()) {
+ mcmp, err := utils.NewMySQLCompare(t, vtParams, mysqlParams)
+ require.NoError(t, err)
+
+ deleteAll := func() {
+ tables := []string{"twopc_user"}
+ for _, table := range tables {
+ _, _ = mcmp.ExecAndIgnore("delete from " + table)
+ }
+ }
+
+ deleteAll()
+
+ return mcmp, func() {
+ deleteAll()
+ mcmp.Close()
+ cluster.PanicHandler(t)
+ }
 }
 
 type extractInterestingValues func(dtidMap map[string]string, vals []sqltypes.Value) []sqltypes.Value
@@ -147,7 +183,8 @@ var tables = map[string]extractInterestingValues{
  },
  "ks.redo_statement": func(dtidMap map[string]string, vals []sqltypes.Value) (out []sqltypes.Value) {
  dtid := getDTID(dtidMap, vals[0].ToString())
- out = append([]sqltypes.Value{sqltypes.NewVarChar(dtid)}, vals[1:]...)
+ stmt := getStatement(vals[2].ToString())
+ out = append([]sqltypes.Value{sqltypes.NewVarChar(dtid)}, vals[1], sqltypes.TestValue(sqltypes.Blob, stmt))
  return
  },
  "ks.twopc_user": func(_ map[string]string, vals []sqltypes.Value) []sqltypes.Value { return vals },
@@ -167,6 +204,28 @@ func getDTID(dtidMap map[string]string, dtKey string) string {
  return dtid
 }
 
+func getStatement(stmt string) string {
+ var sKey string
+ var prefix string
+ switch {
+ case strings.HasPrefix(stmt, "savepoint"):
+ prefix = "savepoint-"
+ sKey = stmt[9:]
+ case strings.HasPrefix(stmt, "rollback to"):
+ prefix = "rollback-"
+ sKey = stmt[11:]
+ default:
+ return stmt
+ }
+
+ sid, exists := sm.stmt[sKey]
+ if !exists {
+ sid = fmt.Sprintf("%d", len(sm.stmt)+1)
+ sm.stmt[sKey] = sid
+ }
+ return prefix + sid
+}
+
 func runVStream(t *testing.T, ctx context.Context, ch chan *binlogdatapb.VEvent, vtgateConn *vtgateconn.VTGateConn) {
  vgtid := &binlogdatapb.VGtid{
  ShardGtids: []*binlogdatapb.ShardGtid{
@@ -272,3 +331,13 @@ func prettyPrint(v interface{}) string {
  }
  return string(b)
 }
+
+type stmtMapper struct {
+ stmt map[string]string
+}
+
+var sm = &stmtMapper{stmt: make(map[string]string)}
+
+func (sm *stmtMapper) reset() {
+ sm.stmt = make(map[string]string)
+}