builtins,workload: add setseed function; use it to make RSW more deterministic

docs/generated/sql/functions.md

@@ -978,6 +978,8 @@ available replica will error.</p>
 </span></td><td>Immutable</td></tr>
 <tr><td><a name="round"></a><code>round(val: <a href="float.html">float</a>) &rarr; <a href="float.html">float</a></code></td><td><span class="funcdesc"><p>Rounds <code>val</code> to the nearest integer using half to even (banker’s) rounding.</p>
 </span></td><td>Immutable</td></tr>
+<tr><td><a name="setseed"></a><code>setseed(seed: <a href="float.html">float</a>) &rarr; void</code></td><td><span class="funcdesc"><p>Sets the seed for subsequent random() calls in this session (value between -1.0 and 1.0, inclusive). There are no guarantees as to how this affects the seed of random() calls that appear in the same query as setseed().</p>
+</span></td><td>Volatile</td></tr>
 <tr><td><a name="sign"></a><code>sign(val: <a href="decimal.html">decimal</a>) &rarr; <a href="decimal.html">decimal</a></code></td><td><span class="funcdesc"><p>Determines the sign of <code>val</code>: <strong>1</strong> for positive; <strong>0</strong> for 0 values; <strong>-1</strong> for negative.</p>
 </span></td><td>Immutable</td></tr>
 <tr><td><a name="sign"></a><code>sign(val: <a href="float.html">float</a>) &rarr; <a href="float.html">float</a></code></td><td><span class="funcdesc"><p>Determines the sign of <code>val</code>: <strong>1</strong> for positive; <strong>0</strong> for 0 values; <strong>-1</strong> for negative.</p>

pkg/sql/conn_executor.go

@@ -81,6 +81,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log/severity"
 	"github.com/cockroachdb/cockroach/pkg/util/metric"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/sentryutil"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
@@ -3650,6 +3651,8 @@ func (ex *connExecutor) initEvalCtx(ctx context.Context, evalCtx *extendedEvalCo
 		indexUsageStats:      ex.indexUsageStats,
 		statementPreparer:    ex,
 	}
+	rng, _ := randutil.NewPseudoRand()
+	evalCtx.RNG = rng
 	evalCtx.copyFromExecCfg(ex.server.cfg)
 }
 

pkg/sql/distsql/BUILD.bazel

@@ -29,6 +29,7 @@ go_library(
         "//pkg/util/log",
         "//pkg/util/mon",
         "//pkg/util/pprofutil",
+        "//pkg/util/randutil",
         "//pkg/util/timeutil",
         "//pkg/util/tochar",
         "//pkg/util/tracing",

pkg/sql/distsql/server.go

@@ -39,6 +39,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
 	"github.com/cockroachdb/cockroach/pkg/util/pprofutil"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tochar"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
@@ -375,6 +376,8 @@ func (ds *ServerImpl) setupFlow(
 			RangeStatsFetcher:         ds.ServerConfig.RangeStatsFetcher,
 			ULIDEntropy:               ulid.Monotonic(crypto_rand.Reader, 0),
 		}
+		rng, _ := randutil.NewPseudoRand()
+		evalCtx.RNG = rng
 		// Most processors will override this Context with their own context in
 		// ProcessorBase. StartInternal().
 		evalCtx.SetDeprecatedContext(ctx)

pkg/sql/logictest/testdata/logic_test/builtin_function

@@ -150,6 +150,31 @@ SELECT random() * 0.0
 ----
 0
 
+statement ok
+SELECT setseed(0.1)
+
+# The previous call to setseed() should affect the next query.
+query IR
+SELECT i, random() FROM ROWS FROM (generate_series(1, 5)) AS i ORDER by i
+----
+1  0.8090535001228529
+2  0.33363615433097443
+3  0.10208231032120892
+4  0.43716650508717286
+5  0.16686635296305902
+
+statement ok
+SELECT setseed(0.1)
+
+query IR
+SELECT i, random() FROM ROWS FROM (generate_series(1, 5)) AS i ORDER by i
+----
+1  0.8090535001228529
+2  0.33363615433097443
+3  0.10208231032120892
+4  0.43716650508717286
+5  0.16686635296305902
+
 # Concatenating 'empty' because the empty string doesn't work in these tests.
 query T
 SELECT concat() || 'empty'

pkg/sql/planner.go

@@ -56,6 +56,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/ulid"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/logtags"
@@ -557,6 +558,8 @@ func internalExtendedEvalCtx(
 		statsProvider:   sqlStatsProvider,
 		jobs:            newTxnJobsCollection(),
 	}
+	rng, _ := randutil.NewPseudoRand()
+	ret.RNG = rng
 	ret.SetDeprecatedContext(ctx)
 	ret.copyFromExecCfg(execCfg)
 	return ret

pkg/sql/schema_changer.go

@@ -66,6 +66,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log/eventpb"
 	"github.com/cockroachdb/cockroach/pkg/util/log/logpb"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/ulid"
@@ -2646,6 +2647,8 @@ func createSchemaChangeEvalCtx(
 			ULIDEntropy:          ulid.Monotonic(crypto_rand.Reader, 0),
 		},
 	}
+	rng, _ := randutil.NewPseudoRand()
+	evalCtx.RNG = rng
 	// TODO(andrei): This is wrong (just like on the main code path on
 	// setupFlow). Each processor should override Ctx with its own context.
 	evalCtx.SetDeprecatedContext(ctx)

pkg/sql/schemachanger/scbuild/BUILD.bazel

@@ -57,6 +57,7 @@ go_library(
         "//pkg/util/log/eventpb",
         "//pkg/util/log/logpb",
         "//pkg/util/mon",
+        "//pkg/util/randutil",
         "//pkg/util/ulid",
         "//pkg/util/uuid",
         "@com_github_cockroachdb_errors//:errors",

pkg/sql/schemachanger/scbuild/tree_context_builder.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/ulid"
 )
 
@@ -64,6 +65,8 @@ func newEvalCtx(ctx context.Context, d Dependencies) *eval.Context {
 		DescIDGenerator:      d.DescIDGenerator(),
 		ULIDEntropy:          ulid.Monotonic(crypto_rand.Reader, 0),
 	}
+	rng, _ := randutil.NewPseudoRand()
+	evalCtx.RNG = rng
 	evalCtx.SetDeprecatedContext(ctx)
 	return evalCtx
 }

pkg/sql/sem/builtins/builtins.go

@@ -24,7 +24,6 @@ import (
 	"hash/fnv"
 	"math"
 	"math/bits"
-	"math/rand"
 	"net"
 	"regexp/syntax"
 	"strings"
@@ -2090,15 +2089,34 @@ var regularBuiltins = map[string]builtinDefinition{
 		tree.Overload{
 			Types:      tree.ParamTypes{},
 			ReturnType: tree.FixedReturnType(types.Float),
-			Fn: func(_ context.Context, _ *eval.Context, args tree.Datums) (tree.Datum, error) {
-				return tree.NewDFloat(tree.DFloat(rand.Float64())), nil
+			Fn: func(_ context.Context, evalCtx *eval.Context, args tree.Datums) (tree.Datum, error) {
+				return tree.NewDFloat(tree.DFloat(evalCtx.RNG.Float64())), nil
 			},
 			Info: "Returns a random floating-point number between 0 (inclusive) and 1 (exclusive). " +
 				"Note that the value contains at most 53 bits of randomness.",
 			Volatility: volatility.Volatile,
 		},
 	),
 
+	"setseed": makeBuiltin(
+		defProps(),
+		tree.Overload{
+			Types:      tree.ParamTypes{{Name: "seed", Typ: types.Float}},
+			ReturnType: tree.FixedReturnType(types.Void),
+			Fn: func(_ context.Context, evalCtx *eval.Context, args tree.Datums) (tree.Datum, error) {
+				seed := tree.MustBeDFloat(args[0])
+				if seed < -1.0 || seed > 1.0 {
+					return nil, pgerror.Newf(pgcode.InvalidParameterValue, "setseed parameter %f is out of allowed range [-1,1]", seed)
+				}
+				evalCtx.RNG.Seed(int64(math.Float64bits(float64(seed))))
+				return tree.DVoidDatum, nil
+			},
+			Info: "Sets the seed for subsequent random() calls in this session (value between -1.0 and 1.0, inclusive). " +
+				"There are no guarantees as to how this affects the seed of random() calls that appear in the same query as setseed().",
+			Volatility: volatility.Volatile,
+		},
+	),
+
 	"unique_rowid": makeBuiltin(
 		tree.FunctionProperties{
 			Category: builtinconstants.CategoryIDGeneration,

pkg/sql/sem/builtins/fixed_oids.go

@@ -2563,6 +2563,7 @@ var builtinOidsArray = []string{
 	2595: `array_agg(arg1: pg_lsn[]) -> pg_lsn[][]`,
 	2596: `array_agg(arg1: refcursor[]) -> refcursor[][]`,
 	2597: `array_agg(arg1: tuple[]) -> tuple[][]`,
+	2598: `setseed(seed: float) -> void`,
 }
 
 var builtinOidsBySignature map[string]oid.Oid

pkg/sql/sem/eval/context.go

@@ -13,6 +13,7 @@ package eval
 import (
 	"context"
 	"math"
+	"math/rand"
 	"time"
 
 	apd "github.com/cockroachdb/apd/v3"
@@ -288,6 +289,9 @@ type Context struct {
 
 	// ULIDEntropy is the entropy source for ULID generation.
 	ULIDEntropy ulid.MonotonicReader
+
+	// RNG is the random number generator use for the "random" built-in function.
+	RNG *rand.Rand
 }
 
 // JobsProfiler is the interface used to fetch job specific execution details

pkg/workload/schemachange/operation_generator.go

@@ -3256,6 +3256,9 @@ func (og *operationGenerator) getTableColumns(
 func (og *operationGenerator) randColumn(
 	ctx context.Context, tx pgx.Tx, tableName tree.TableName, pctExisting int,
 ) (string, error) {
+	if err := og.setSeedInDB(ctx, tx); err != nil {
+		return "", err
+	}
 	if og.randIntn(100) >= pctExisting {
 		// We make a unique name for all columns by prefixing them with the table
 		// index to make it easier to reference columns from different tables.
@@ -3282,6 +3285,9 @@ ORDER BY random()
 func (og *operationGenerator) randColumnWithMeta(
 	ctx context.Context, tx pgx.Tx, tableName tree.TableName, pctExisting int,
 ) (column, error) {
+	if err := og.setSeedInDB(ctx, tx); err != nil {
+		return column{}, err
+	}
 	if og.randIntn(100) >= pctExisting {
 		// We make a unique name for all columns by prefixing them with the table
 		// index to make it easier to reference columns from different tables.
@@ -3371,7 +3377,9 @@ func (og *operationGenerator) randChildColumnForFkRelation(
 func (og *operationGenerator) randParentColumnForFkRelation(
 	ctx context.Context, tx pgx.Tx, unique bool,
 ) (*tree.TableName, *column, error) {
-
+	if err := og.setSeedInDB(ctx, tx); err != nil {
+		return nil, nil, err
+	}
 	subQuery := strings.Builder{}
 	subQuery.WriteString(`
 		SELECT table_schema, table_name, column_name, crdb_sql_type, is_nullable, contype, conkey
@@ -3444,6 +3452,9 @@ func (og *operationGenerator) randParentColumnForFkRelation(
 func (og *operationGenerator) randConstraint(
 	ctx context.Context, tx pgx.Tx, tableName string,
 ) (string, error) {
+	if err := og.setSeedInDB(ctx, tx); err != nil {
+		return "", err
+	}
 	q := fmt.Sprintf(`
   SELECT constraint_name
     FROM [SHOW CONSTRAINTS FROM %s]
@@ -3461,6 +3472,9 @@ ORDER BY random()
 func (og *operationGenerator) randIndex(
 	ctx context.Context, tx pgx.Tx, tableName tree.TableName, pctExisting int,
 ) (string, error) {
+	if err := og.setSeedInDB(ctx, tx); err != nil {
+		return "", err
+	}
 	if og.randIntn(100) >= pctExisting {
 		// We make a unique name for all indices by prefixing them with the table
 		// index to make it easier to reference columns from different tables.
@@ -3485,7 +3499,9 @@ ORDER BY random()
 func (og *operationGenerator) randSequence(
 	ctx context.Context, tx pgx.Tx, pctExisting int, desiredSchema string,
 ) (*tree.TableName, error) {
-
+	if err := og.setSeedInDB(ctx, tx); err != nil {
+		return nil, err
+	}
 	if desiredSchema != "" {
 		if og.randIntn(100) >= pctExisting {
 			treeSeqName := tree.MakeTableNameFromPrefix(tree.ObjectNamePrefix{
@@ -3587,7 +3603,9 @@ ORDER BY random()
 func (og *operationGenerator) randTable(
 	ctx context.Context, tx pgx.Tx, pctExisting int, desiredSchema string,
 ) (*tree.TableName, error) {
-
+	if err := og.setSeedInDB(ctx, tx); err != nil {
+		return nil, err
+	}
 	if desiredSchema != "" {
 		if og.randIntn(100) >= pctExisting {
 			treeTableName := tree.MakeTableNameFromPrefix(tree.ObjectNamePrefix{
@@ -3658,6 +3676,9 @@ ORDER BY random()
 func (og *operationGenerator) randView(
 	ctx context.Context, tx pgx.Tx, pctExisting int, desiredSchema string,
 ) (*tree.TableName, error) {
+	if err := og.setSeedInDB(ctx, tx); err != nil {
+		return nil, err
+	}
 	if desiredSchema != "" {
 		if og.randIntn(100) >= pctExisting {
 			treeViewName := tree.MakeTableNameFromPrefix(tree.ObjectNamePrefix{
@@ -3702,6 +3723,9 @@ func (og *operationGenerator) randView(
 		}, tree.Name(fmt.Sprintf("view_%s", og.newUniqueSeqNumSuffix())))
 		return &treeViewName, nil
 	}
+	if err := og.setSeedInDB(ctx, tx); err != nil {
+		return nil, err
+	}
 	const q = `
   SELECT schema_name, table_name
     FROM [SHOW TABLES]
@@ -3812,6 +3836,9 @@ func (og *operationGenerator) createSchema(ctx context.Context, tx pgx.Tx) (*opS
 func (og *operationGenerator) randSchema(
 	ctx context.Context, tx pgx.Tx, pctExisting int,
 ) (string, error) {
+	if err := og.setSeedInDB(ctx, tx); err != nil {
+		return "", err
+	}
 	if og.randIntn(100) >= pctExisting {
 		return fmt.Sprintf("schema_%s", og.newUniqueSeqNumSuffix()), nil
 	}
@@ -4309,6 +4336,11 @@ func (og *operationGenerator) randIntn(topBound int) int {
 	return og.params.rng.Intn(topBound)
 }
 
+// randFloat64 returns an float64 in the range [0,1.0).
+func (og *operationGenerator) randFloat64() float64 {
+	return og.params.rng.Float64()
+}
+
 func (og *operationGenerator) newUniqueSeqNumSuffix() string {
 	og.params.seqNum++
 	return fmt.Sprintf("w%d_%d", og.params.workerID, og.params.seqNum)
@@ -4335,7 +4367,7 @@ func (og *operationGenerator) typeFromTypeName(
 }
 
 // Check if the test is running with a mixed version cluster, with a version
-// less than or equal to the target version number. This can be used to detect
+// less than the target version number. This can be used to detect
 // in mixed version environments if certain errors should be encountered.
 func isClusterVersionLessThan(
 	ctx context.Context, tx pgx.Tx, targetVersion roachpb.Version,
@@ -4349,5 +4381,19 @@ func isClusterVersionLessThan(
 	if err != nil {
 		return false, err
 	}
-	return clusterVersion.LessEq(targetVersion), nil
+	return clusterVersion.Less(targetVersion), nil
+}
+
+func (og *operationGenerator) setSeedInDB(ctx context.Context, tx pgx.Tx) error {
+	if notSupported, err := isClusterVersionLessThan(ctx, tx, clusterversion.V24_1.Version()); err != nil {
+		return err
+	} else if notSupported {
+		// To allow the schemachange workload to work in a mixed-version state,
+		// this should not be treated as an error.
+		return nil
+	}
+	if _, err := tx.Exec(ctx, "SELECT setseed($1)", og.randFloat64()); err != nil {
+		return err
+	}
+	return nil
 }

pkg/workload/schemachange/schemachange.go

@@ -201,19 +201,17 @@ func (s *schemaChange) Ops(
 		return workload.QueryLoad{}, err
 	}
 	stdoutLog := makeAtomicLog(os.Stdout)
-	rng, seed := randutil.NewTestRand()
+	_, seed := randutil.NewTestRand()
 	stdoutLog.printLn(fmt.Sprintf("using random seed: %d", seed))
-	ops := newDeck(rng, opWeights...)
-	// A separate deck is constructed of only schema changes supported
-	// by the declarative schema changer. This deck has equal weights,
-	// only for supported schema changes.
+	// A separate weighting is constructed of only schema changes supported by the
+	// declarative schema changer. This will be used to make a per-worker deck
+	// that has equal weights, only for supported schema changes.
 	declarativeOpWeights := make([]int, len(opWeights))
 	for idx, weight := range opWeights {
 		if _, ok := opDeclarativeVersion[opType(idx)]; ok {
 			declarativeOpWeights[idx] = weight
 		}
 	}
-	declarativeOps := newDeck(rng, declarativeOpWeights...)
 
 	ql := workload.QueryLoad{
 		SQLDatabase: sqlDatabase,
@@ -251,12 +249,15 @@ func (s *schemaChange) Ops(
 		// operations.
 		workerRng := randutil.NewTestRandWithSeed(seed + int64(i))
 
-		// Each worker needs its own sequence number generator so that the names of
-		// generated objects are deterministic across runs.
+		// Each worker needs its own sequence number generator and operation deck so
+		// that the names of generated objects and operations are deterministic
+		// across runs.
 		seqNum, err := s.initSeqNum(ctx, pool, i)
 		if err != nil {
 			return workload.QueryLoad{}, err
 		}
+		ops := newDeck(workerRng, opWeights...)
+		declarativeOps := newDeck(workerRng, declarativeOpWeights...)
 
 		opGeneratorParams := operationGeneratorParams{
 			workerID:           i,