sql/schemachanger: add sampling for pause / rollback tests

Previously, the pause and rollback tests for the declarative schema
changer could timeout due to the number of stages being tested. To
address this, this patch adds supports for sampling stages in plans for
rollback / pause testing.

Fixes: cockroachdb#155668
Fixes: cockroachdb#155502

Release note: None

Author: fqazi

pkg/sql/schemachanger/sctest/backup.go

@@ -52,7 +52,8 @@ func BackupSuccess(t *testing.T, path string, factory TestServerFactory) {
 		backupArgs = backupSuccessPrepare(t, factory, spec, dbName)
 	}, func(t *testing.T, cs CumulativeTestCaseSpec) {
 		backupSuccess(t, backupArgs, cs)
-	})
+	},
+		nil /*samplingFn */)
 }
 
 // BackupRollbacks tests that the schema changer can handle being backed up
@@ -70,7 +71,8 @@ func BackupRollbacks(t *testing.T, path string, factory TestServerFactory) {
 	factory = factory.WithSchemaLockDisabled()
 	cumulativeTestForEachPostCommitStage(t, path, factory, nil, func(t *testing.T, cs CumulativeTestCaseSpec) {
 		backupRollbacks(t, factory, cs)
-	})
+	},
+		nil /* samplingFn */)
 }
 
 // BackupSuccessMixedVersion is like BackupSuccess but in a mixed-version
@@ -96,7 +98,8 @@ func BackupSuccessMixedVersion(t *testing.T, path string, factory TestServerFact
 		backupArgs = backupSuccessPrepare(t, factory, spec, dbName)
 	}, func(t *testing.T, cs CumulativeTestCaseSpec) {
 		backupSuccess(t, backupArgs, cs)
-	})
+	},
+		nil /* samplingFn */)
 }
 
 // BackupRollbacksMixedVersion is like BackupRollbacks but in a mixed-version
@@ -116,7 +119,7 @@ func BackupRollbacksMixedVersion(t *testing.T, path string, factory TestServerFa
 	factory = factory.WithSchemaLockDisabled()
 	cumulativeTestForEachPostCommitStage(t, path, factory, nil, func(t *testing.T, cs CumulativeTestCaseSpec) {
 		backupRollbacks(t, factory, cs)
-	})
+	}, nil /* samplingFn */)
 }
 
 // runAllBackups runs all the backup tests, disabling the random skipping.

pkg/sql/schemachanger/sctest/cumulative.go

@@ -117,7 +117,7 @@ func Rollback(t *testing.T, relPath string, factory TestServerFactory) {
 		}
 		factory.WithSchemaChangerKnobs(knobs).Run(ctx, t, runfn)
 	}
-	cumulativeTestForEachPostCommitStage(t, relPath, factory, nil, testRollbackCase)
+	cumulativeTestForEachPostCommitStage(t, relPath, factory, nil, testRollbackCase, sampleAllPostCommitRevertible /* enableSampling */)
 }
 
 // ExecuteWithDMLInjection tests that the schema changer behaviour is sane
@@ -335,7 +335,8 @@ func Pause(t *testing.T, path string, factory TestServerFactory) {
 
 	cumulativeTestForEachPostCommitStage(t, path, factory, nil, func(t *testing.T, cs CumulativeTestCaseSpec) {
 		pause(t, factory, cs)
-	})
+	},
+		sampleAllPostCommitStages /* enableSampling */)
 }
 
 // PauseMixedVersion is like Pause but in a mixed-version cluster which gets
@@ -349,7 +350,8 @@ func PauseMixedVersion(t *testing.T, path string, factory TestServerFactory) {
 	factory.WithMixedVersion()
 	cumulativeTestForEachPostCommitStage(t, path, factory, nil, func(t *testing.T, cs CumulativeTestCaseSpec) {
 		pause(t, factory, cs)
-	})
+	},
+		sampleAllPostCommitStages /* enableSampling */)
 }
 
 func pause(t *testing.T, factory TestServerFactory, cs CumulativeTestCaseSpec) {

pkg/sql/schemachanger/sctest/framework.go

@@ -8,8 +8,10 @@ package sctest
 import (
 	"context"
 	gosql "database/sql"
+	"flag"
 	"fmt"
 	"math"
+	"math/rand"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -634,6 +636,31 @@ type CumulativeTestCaseSpec struct {
 	CreateDatabaseStmt string
 }
 
+// sampleAllPostCommitRevertible samples all post commit revertible stages, and
+// limits testing to maxStagesToTest if *runAllCumulative is not set.
+func sampleAllPostCommitRevertible(testCases []CumulativeTestCaseSpec) []CumulativeTestCaseSpec {
+	newTestCases := make([]CumulativeTestCaseSpec, 0, len(testCases))
+	for _, tc := range testCases {
+		if tc.Phase != scop.PostCommitNonRevertiblePhase {
+			newTestCases = append(newTestCases, tc)
+		}
+	}
+	return sampleAllPostCommitStages(newTestCases)
+}
+
+// sampleAllPostCommitStages samples all post-commit stages, and limits these to
+// maxStagesToTest if *runAllCumulative is not set.
+func sampleAllPostCommitStages(testCases []CumulativeTestCaseSpec) []CumulativeTestCaseSpec {
+	if len(testCases) > maxStagesToTest && !(*runAllCumulative) {
+		// Shuffle and pick up to maxStagesToTest.
+		rand.Shuffle(len(testCases), func(i, j int) {
+			testCases[i], testCases[j] = testCases[j], testCases[i]
+		})
+		testCases = testCases[:maxStagesToTest]
+	}
+	return testCases
+}
+
 func (cs CumulativeTestCaseSpec) run(t *testing.T, fn func(t *testing.T)) bool {
 	var prefix string
 	switch cs.Phase {
@@ -647,6 +674,12 @@ func (cs CumulativeTestCaseSpec) run(t *testing.T, fn func(t *testing.T)) bool {
 	return t.Run(fmt.Sprintf("%s_stage_%d_of_%d", prefix, cs.StageOrdinal, cs.StagesCount), fn)
 }
 
+// / runAllCumulative used to disable sampling for cumulative tests.
+var runAllCumulative = flag.Bool(
+	"run-all-cumulative", false,
+	"if true, run all cumulative instead of a random subset",
+)
+
 // cumulativeTestForEachPostCommitStage invokes `tf` once for each stage in the
 // PostCommitPhase.
 func cumulativeTestForEachPostCommitStage(
@@ -655,6 +688,7 @@ func cumulativeTestForEachPostCommitStage(
 	factory TestServerFactory,
 	prepFn func(t *testing.T, spec CumulativeTestSpec, dbName string),
 	tf func(t *testing.T, spec CumulativeTestCaseSpec),
+	samplingFn func([]CumulativeTestCaseSpec) []CumulativeTestCaseSpec,
 ) {
 	testFunc := func(t *testing.T, spec CumulativeTestSpec) {
 		// Skip this test if any of the stmts is not fully supported.
@@ -720,6 +754,10 @@ func cumulativeTestForEachPostCommitStage(
 		if prepFn != nil {
 			prepFn(t, spec, dbName)
 		}
+		// If sampling is enabled limit the number of stages executed.
+		if samplingFn != nil {
+			testCases = samplingFn(testCases)
+		}
 		for _, tc := range testCases {
 			fn := func(t *testing.T) {
 				tf(t, tc)