CI: Address our two flakiest tests (#19587)

Signed-off-by: Matt Lord <mattalord@gmail.com>

Author: mattlord

go/vt/vttablet/tabletmanager/framework_test.go

@@ -23,6 +23,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -59,6 +60,8 @@ const (
 	gtidPosition = "16b1039f-22b6-11ed-b765-0a43f95f28a3:1-220"
 )
 
+var testEnvCounter atomic.Int64
+
 func init() {
 	tabletconn.RegisterDialer("grpc", func(ctx context.Context, tablet *topodatapb.Tablet, failFast grpcclient.FailFast) (queryservice.QueryService, error) {
 		return &tabletconntest.FakeQueryService{
@@ -90,12 +93,13 @@ type testEnv struct {
 
 func newTestEnv(t *testing.T, ctx context.Context, sourceKeyspace string, sourceShards []string) *testEnv {
 	vttablet.InitVReplicationConfigDefaults()
+	protoName := fmt.Sprintf("%s_%d", t.Name(), testEnvCounter.Add(1))
 	tenv := &testEnv{
 		ctx:       context.Background(),
 		tmc:       newFakeTMClient(),
 		cells:     []string{"zone1"},
 		dbName:    "tmtestdb",
-		protoName: t.Name(),
+		protoName: protoName,
 	}
 	tenv.mu.Lock()
 	defer tenv.mu.Unlock()
@@ -104,19 +108,19 @@ func newTestEnv(t *testing.T, ctx context.Context, sourceKeyspace string, source
 	tenv.tmc.sourceShards = sourceShards
 	tenv.tmc.schema = defaultSchema
 
-	tabletconn.RegisterDialer(t.Name(), func(ctx context.Context, tablet *topodatapb.Tablet, failFast grpcclient.FailFast) (queryservice.QueryService, error) {
+	tabletconn.RegisterDialer(protoName, func(ctx context.Context, tablet *topodatapb.Tablet, failFast grpcclient.FailFast) (queryservice.QueryService, error) {
 		tenv.mu.Lock()
 		defer tenv.mu.Unlock()
 		if qs, ok := tenv.tmc.tablets[int(tablet.Alias.Uid)]; ok {
 			return qs, nil
 		}
 		return nil, fmt.Errorf("tablet %d not found", tablet.Alias.Uid)
 	})
-	tabletconntest.SetProtocol("go.vt.vttablet.tabletmanager.framework_test_"+t.Name(), tenv.protoName)
-	tmclient.RegisterTabletManagerClientFactory(t.Name(), func() tmclient.TabletManagerClient {
+	tabletconntest.SetProtocol("go.vt.vttablet.tabletmanager.framework_test_"+protoName, protoName)
+	tmclient.RegisterTabletManagerClientFactory(protoName, func() tmclient.TabletManagerClient {
 		return tenv.tmc
 	})
-	tmclienttest.SetProtocol("go.vt.vttablet.tabletmanager.framework_test_"+t.Name(), tenv.protoName)
+	tmclienttest.SetProtocol("go.vt.vttablet.tabletmanager.framework_test_"+protoName, protoName)
 
 	tenv.db = fakesqldb.New(t)
 	tenv.mysqld = mysqlctl.NewFakeMysqlDaemon(tenv.db)

go/vt/vttablet/tabletmanager/rpc_vreplication_test.go

@@ -1840,7 +1840,7 @@ func addInvariants(dbClient *binlogplayer.MockDBClient, vreplID, sourceTabletUID
 		"0",
 	))
 	dbClient.AddInvariant(fmt.Sprintf(updatePickedSourceTablet, cell, sourceTabletUID, vreplID), &sqltypes.Result{})
-	dbClient.AddInvariant("update _vt.vreplication set state='Running', message='' where id=1", &sqltypes.Result{})
+	dbClient.AddInvariant("update _vt.vreplication set state='Running', message=left('', 1000) where id=1", &sqltypes.Result{})
 	dbClient.AddInvariant(vreplication.SqlMaxAllowedPacket, sqltypes.MakeTestResult(
 		sqltypes.MakeTestFields(
 			"max_allowed_packet",

go/vt/vttablet/tabletmanager/vreplication/framework_test.go

@@ -74,6 +74,7 @@ var (
 	vrepldb                  = "vrepl"
 	globalDBQueries          = make(chan string, 1000)
 	lastMultiExecQuery       = ""
+	lastMultiExecQueryMu     sync.Mutex
 	testForeignKeyQueries    = false
 	testSetForeignKeyQueries = false
 	doNotLogDBQueries        = false
@@ -608,10 +609,18 @@ func (dbc *realDBClient) ExecuteFetchMulti(query string, maxrows int) ([]*sqltyp
 		}
 		results = append(results, qr)
 	}
+	lastMultiExecQueryMu.Lock()
 	lastMultiExecQuery = query
+	lastMultiExecQueryMu.Unlock()
 	return results, nil
 }
 
+func getLastMultiExecQuery() string {
+	lastMultiExecQueryMu.Lock()
+	defer lastMultiExecQueryMu.Unlock()
+	return lastMultiExecQuery
+}
+
 func (dbc *realDBClient) SupportsCapability(capability capabilities.FlavorCapability) (bool, error) {
 	return dbc.conn.SupportsCapability(capability)
 }

go/vt/vttablet/tabletmanager/vreplication/vplayer_flaky_test.go

@@ -3810,7 +3810,6 @@ func TestPlayerBatchMode(t *testing.T) {
 				require.LessOrEqual(t, len(stmt), maxBatchSize, "expected output statement is longer than the max batch size (%d): %s", maxBatchSize, stmt)
 			}
 			expectNontxQueries(t, output, recvTimeout)
-			time.Sleep(1 * time.Second)
 			if tcase.table != "" {
 				expectData(t, tcase.table, tcase.data)
 			}
@@ -3822,15 +3821,35 @@ func TestPlayerBatchMode(t *testing.T) {
 			expectedBulkInserts += tcase.expectedBulkInserts
 			expectedTrxBatchCommits++ // Should only ever be 1 per test case
 			expectedTrxBatchExecs += tcase.expectedNonCommitBatches
-			if tcase.expectedInLastBatch != "" { // We expect the trx to be split
-				require.Regexpf(t, regexp.MustCompile(fmt.Sprintf(trxLastBatchExpectRE, regexp.QuoteMeta(tcase.expectedInLastBatch))), lastMultiExecQuery, "Unexpected batch statement: %s", lastMultiExecQuery)
+
+			// Poll until the batch query and stats counters are updated.
+			// These are set asynchronously on the vplayer goroutine after
+			// the commit completes, so we poll rather than using a fixed sleep.
+			var batchRE *regexp.Regexp
+			if tcase.expectedInLastBatch != "" {
+				batchRE = regexp.MustCompile(fmt.Sprintf(trxLastBatchExpectRE, regexp.QuoteMeta(tcase.expectedInLastBatch)))
 			} else {
-				require.Regexpf(t, regexp.MustCompile(fmt.Sprintf(trxFullBatchExpectRE, regexp.QuoteMeta(strings.Join(tcase.output, ";")))), lastMultiExecQuery, "Unexpected batch statement: %s", lastMultiExecQuery)
+				batchRE = regexp.MustCompile(fmt.Sprintf(trxFullBatchExpectRE, regexp.QuoteMeta(strings.Join(tcase.output, ";"))))
 			}
-			require.Equal(t, expectedBulkInserts, stats.BulkQueryCount.Counts()["insert"], "expected %d bulk inserts but got %d", expectedBulkInserts, stats.BulkQueryCount.Counts()["insert"])
-			require.Equal(t, expectedBulkDeletes, stats.BulkQueryCount.Counts()["delete"], "expected %d bulk deletes but got %d", expectedBulkDeletes, stats.BulkQueryCount.Counts()["delete"])
-			require.Equal(t, expectedTrxBatchExecs, stats.TrxQueryBatchCount.Counts()["without_commit"], "expected %d trx batch execs but got %d", expectedTrxBatchExecs, stats.TrxQueryBatchCount.Counts()["without_commit"])
-			require.Equal(t, expectedTrxBatchCommits, stats.TrxQueryBatchCount.Counts()["with_commit"], "expected %d trx batch commits but got %d", expectedTrxBatchCommits, stats.TrxQueryBatchCount.Counts()["with_commit"])
+			require.Eventually(t, func() bool {
+				got := getLastMultiExecQuery()
+				if !batchRE.MatchString(got) {
+					return false
+				}
+				if stats.BulkQueryCount.Counts()["insert"] != expectedBulkInserts {
+					return false
+				}
+				if stats.BulkQueryCount.Counts()["delete"] != expectedBulkDeletes {
+					return false
+				}
+				if stats.TrxQueryBatchCount.Counts()["without_commit"] != expectedTrxBatchExecs {
+					return false
+				}
+				if stats.TrxQueryBatchCount.Counts()["with_commit"] != expectedTrxBatchCommits {
+					return false
+				}
+				return true
+			}, 10*time.Second, 100*time.Millisecond, "batch query or stats mismatch after timeout; lastMultiExecQuery: %s", getLastMultiExecQuery())
 		})
 	}
 }