metamorphic: add environment-based overrides

This allows developers to override specific metamorphic variables:

   ./dev test ./pkg/foo -- --test_env COCKROACH_INTERNAL_METAMORPHIC_OVERRIDES=var1=val1,var2=val2

Metamorphic values not specified are still randomly chosen.

In a future commit, it would also be nice to be able to specify a file
that allows you to specify all values.

Epic: none
Informs #114039

pkg/util/metamorphic/BUILD.bazel

@@ -21,7 +21,14 @@ go_library(
 
 go_test(
     name = "metamorphic_test",
-    srcs = ["constants_test.go"],
+    srcs = [
+        "constants_test.go",
+        "main_test.go",
+    ],
     embed = [":metamorphic"],
-    deps = ["@com_github_stretchr_testify//require"],
+    deps = [
+        "//pkg/util/log",
+        "//pkg/util/randutil",
+        "@com_github_stretchr_testify//require",
+    ],
 )

pkg/util/metamorphic/constants.go

@@ -9,9 +9,12 @@ import (
 	"fmt"
 	"math/rand"
 	"os"
+	"strconv"
+	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/build/bazel"
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
+	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/metamorphic/metamorphicutil"
 	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
@@ -61,52 +64,14 @@ const (
 //
 // The given name is used for logging.
 func ConstantWithTestValue(name string, defaultValue, metamorphicValue int) int {
-	if metamorphicutil.IsMetamorphicBuild {
-		rng.Lock()
-		defer rng.Unlock()
-		if rng.r.Float64() < metamorphicValueProbability {
-			logMetamorphicValue(name, metamorphicValue)
+	return getConstantInternal(name,
+		defaultValue,
+		strconv.Atoi,
+		metamorphicValueProbability,
+		func(r *rand.Rand) int {
 			return metamorphicValue
-		}
-	}
-	return defaultValue
-}
-
-// rng is initialized to a rand.Rand if crdbTestBuild is enabled.
-var rng struct {
-	r *rand.Rand
-	syncutil.Mutex
-}
-
-// DisableMetamorphicEnvVar can be used to disable metamorphic tests for
-// sub-processes. If it exists and is set to something truthy as defined by
-// strconv.ParseBool then metamorphic testing will not be enabled.
-const DisableMetamorphicEnvVar = "COCKROACH_INTERNAL_DISABLE_METAMORPHIC_TESTING"
-
-// Returns true iff the current process is eligible to enable metamorphic
-// variables. When run under Bazel, checking if we are in the Go test wrapper
-// ensures that metamorphic variables are not initialized and logged twice
-// from both the wrapper and the main test process, as both will perform
-// initialization of the test module and its dependencies.
-func metamorphicEligible() bool {
-	if !buildutil.CrdbTestBuild {
-		return false
-	}
-
-	if bazel.InTestWrapper() {
-		return false
-	}
-
-	return true
-}
-
-func init() {
-	if metamorphicEligible() {
-		if !disableMetamorphicTesting {
-			rng.r, _ = randutil.NewTestRand()
-			metamorphicutil.IsMetamorphicBuild = rng.r.Float64() < IsMetamorphicBuildProbability
-		}
-	}
+		},
+		true)
 }
 
 // ConstantWithTestRange is like ConstantWithTestValue except instead of
@@ -115,19 +80,18 @@ func init() {
 //
 // The given name is used for logging.
 func ConstantWithTestRange(name string, defaultValue, min, max int) int {
-	if metamorphicutil.IsMetamorphicBuild {
-		rng.Lock()
-		defer rng.Unlock()
-		if rng.r.Float64() < metamorphicValueProbability {
+	return getConstantInternal(name,
+		defaultValue,
+		strconv.Atoi,
+		metamorphicValueProbability,
+		func(r *rand.Rand) int {
 			ret := min
 			if max > min {
-				ret = int(rng.r.Int31())%(max-min) + min
+				ret = int(r.Int31())%(max-min) + min
 			}
-			logMetamorphicValue(name, ret)
 			return ret
-		}
-	}
-	return defaultValue
+		},
+		true)
 }
 
 // ConstantWithTestBool is like ConstantWithTestValue except it returns the
@@ -138,21 +102,6 @@ func ConstantWithTestBool(name string, defaultValue bool) bool {
 	return constantWithTestBoolInternal(name, defaultValue, true /* doLog */)
 }
 
-func constantWithTestBoolInternal(name string, defaultValue bool, doLog bool) bool {
-	if metamorphicutil.IsMetamorphicBuild {
-		rng.Lock()
-		defer rng.Unlock()
-		if rng.r.Float64() < metamorphicBoolProbability {
-			ret := !defaultValue
-			if doLog {
-				logMetamorphicValue(name, ret)
-			}
-			return ret
-		}
-	}
-	return defaultValue
-}
-
 // ConstantWithTestBoolWithoutLogging is like ConstantWithTestBool except it
 // does not log the value. This is necessary to work around this issue:
 // https://github.com/cockroachdb/cockroach/issues/106667
@@ -161,23 +110,181 @@ func ConstantWithTestBoolWithoutLogging(name string, defaultValue bool) bool {
 	return constantWithTestBoolInternal(name, defaultValue, false /* doLog */)
 }
 
+func constantWithTestBoolInternal(name string, defaultValue bool, doLog bool) bool {
+	return getConstantInternal(name,
+		defaultValue,
+		strconv.ParseBool,
+		metamorphicBoolProbability,
+		func(*rand.Rand) bool { return !defaultValue },
+		doLog)
+}
+
 // ConstantWithTestChoice is like ConstantWithTestValue except it returns a
 // random choice (equally weighted) of the given values. The default value is
 // included in the random choice.
 //
 // The given name is used for logging.
 func ConstantWithTestChoice[T any](name string, defaultValue T, otherValues ...T) T {
-	if metamorphicutil.IsMetamorphicBuild {
-		values := append([]T{defaultValue}, otherValues...)
+	return getConstantInternal(name,
+		defaultValue,
+		func(s string) (T, error) {
+			v, err := parseChoice[T](s)
+			return v.(T), err
+		},
+		1.0,
+		func(r *rand.Rand) T {
+			values := append([]T{defaultValue}, otherValues...)
+			return values[rng.r.Int63n(int64(len(values)))]
+		},
+		true)
+}
+
+// getConstantInternal returns a value of type T for the given name. If the name
+// has been specified as an override, the parseValue func is used to parse the
+// user-provided value. If no such override is specified, it either returns the
+// defaultValue or a value generated by generateValue()
+func getConstantInternal[T any](
+	name string,
+	defaultValue T,
+	parseValue func(string) (T, error),
+	probability float64,
+	generateValue func(r *rand.Rand) T,
+	doLog bool,
+) T {
+	overrideVal, hasOverride := valueFromOverride(name, parseValue)
+	// This is structured so that we make the same number uses of
+	// `rng.r` in the case of a metamoprhic build so that the seed
+	// value is more useful.
+	if metamorphicutil.IsMetamorphicBuild || hasOverride {
 		rng.Lock()
 		defer rng.Unlock()
-		value := values[rng.r.Int63n(int64(len(values)))]
-		logMetamorphicValue(name, value)
-		return value
+		shouldGenerate := rng.r.Float64() < probability
+		if shouldGenerate || hasOverride {
+			ret := generateValue(rng.r)
+			src := randomVal
+			if hasOverride {
+				src = override
+				ret = overrideVal
+			}
+
+			if doLog {
+				logMetamorphicValue(name, ret, src)
+			}
+			return ret
+		}
 	}
 	return defaultValue
 }
 
-func logMetamorphicValue(name string, value interface{}) {
-	fmt.Fprintf(os.Stderr, "initialized metamorphic constant %q with value %v\n", name, value)
+// parseChoice tries to parse the given string into the type T. The returned
+// value should be safe to cast to type T.
+func parseChoice[T any](s string) (any, error) {
+	var zero T
+	switch any(zero).(type) {
+	case string:
+		return s, nil
+	case int, int16, int32, int64:
+		return strconv.Atoi(s)
+	case float32:
+		return strconv.ParseFloat(s, 32)
+	case float64:
+		return strconv.ParseFloat(s, 64)
+	default:
+		panic(fmt.Sprintf("unable to parse %T", zero))
+	}
+}
+
+type metamorphicSource string
+
+const (
+	randomVal metamorphicSource = ""
+	override  metamorphicSource = " (from override)"
+)
+
+func logMetamorphicValue(name string, value interface{}, src metamorphicSource) {
+	logf("initialized metamorphic constant %q with value %v%s\n", name, value, src)
+}
+
+func logf(fmtStr string, args ...interface{}) {
+	fmt.Fprintf(os.Stderr, fmtStr, args...)
+}
+
+const (
+	// DisableMetamorphicEnvVar can be used to disable metamorphic tests for
+	// sub-processes. If it exists and is set to something truthy as defined by
+	// strconv.ParseBool then metamorphic testing will not be enabled.
+	DisableMetamorphicEnvVar = "COCKROACH_INTERNAL_DISABLE_METAMORPHIC_TESTING"
+	// MetamorphicOverridesEnvVar can be used to set metamorphic
+	// variables to specific values.
+	MetamorphicOverridesEnvVar = "COCKROACH_INTERNAL_METAMORPHIC_OVERRIDES"
+)
+
+var (
+	// overrides holds user-provided values parsed from
+	// MetamorphicOverridesEnvVar.
+	overrides = map[string]string{}
+
+	// rng is initialized to a rand.Rand if crdbTestBuild is enabled.
+	rng struct {
+		r *rand.Rand
+		syncutil.Mutex
+	}
+)
+
+// Returns true iff the current process is eligible to enable metamorphic
+// variables. When run under Bazel, checking if we are in the Go test wrapper
+// ensures that metamorphic variables are not initialized and logged twice
+// from both the wrapper and the main test process, as both will perform
+// initialization of the test module and its dependencies.
+func metamorphicEligible() bool {
+	if !buildutil.CrdbTestBuild {
+		return false
+	}
+
+	if bazel.InTestWrapper() {
+		return false
+	}
+
+	return true
+}
+
+// valueFromOverride retuns any user-provided override value for the given name.
+// The passed parser function is used to parse the value from a string
+// representation.
+func valueFromOverride[T any](name string, parser func(string) (T, error)) (T, bool) {
+	if v, ok := overrides[name]; ok {
+		ret, err := parser(v)
+		if err != nil {
+			panic(fmt.Sprintf("malformed value for %s: %s: %s", name, v, err.Error()))
+		}
+		return ret, true
+	} else {
+		var zero T
+		return zero, false
+	}
+}
+
+func init() {
+	if metamorphicEligible() {
+		if !disableMetamorphicTesting {
+			rng.r, _ = randutil.NewTestRand()
+			metamorphicutil.IsMetamorphicBuild = rng.r.Float64() < IsMetamorphicBuildProbability
+		}
+
+		if overrideList, ok := envutil.EnvString(MetamorphicOverridesEnvVar, 0); ok {
+			logf("metamorphic: using overrides from environment\n")
+			setOverridesFromString(overrideList)
+		}
+	}
+}
+
+func setOverridesFromString(overrideList string) {
+	overrideItems := strings.Split(overrideList, ",")
+	for _, item := range overrideItems {
+		itemParts := strings.SplitN(item, "=", 2)
+		if len(itemParts) < 2 {
+			panic(fmt.Sprintf("malformed override: %s", item))
+		}
+		overrides[itemParts[0]] = itemParts[1]
+	}
 }

pkg/util/metamorphic/constants_test.go

@@ -16,3 +16,14 @@ import (
 func TestMetamorphicEligible(t *testing.T) {
 	require.True(t, metamorphicEligible())
 }
+
+// TestMetamorphicFromOverride checks that overrides are used.
+func TestMetamorphicFromOverride(t *testing.T) {
+	setOverridesFromString("val2=7")
+	var (
+		_ = ConstantWithTestRange("val1", 1, 1, 100)
+		v = ConstantWithTestRange("val2", 2, 1, 100)
+		_ = ConstantWithTestRange("val3", 3, 1, 100)
+	)
+	require.Equal(t, 7, v)
+}

pkg/util/metamorphic/main_test.go

@@ -0,0 +1,19 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package metamorphic_test
+
+import (
+	"os"
+	"testing"
+
+	_ "github.com/cockroachdb/cockroach/pkg/util/log" // for flags
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+)
+
+func TestMain(m *testing.M) {
+	randutil.SeedForTests()
+	os.Exit(m.Run())
+}