diff --git a/pkg/kv/kvserver/kvflowcontrol/BUILD.bazel b/pkg/kv/kvserver/kvflowcontrol/BUILD.bazel index 627e9c427c28..077a49e8dd87 100644 --- a/pkg/kv/kvserver/kvflowcontrol/BUILD.bazel +++ b/pkg/kv/kvserver/kvflowcontrol/BUILD.bazel @@ -18,7 +18,7 @@ go_library( "//pkg/settings/cluster", "//pkg/util/admission/admissionpb", "//pkg/util/ctxutil", + "//pkg/util/humanizeutil", "@com_github_cockroachdb_redact//:redact", - "@com_github_dustin_go_humanize//:go-humanize", ], ) diff --git a/pkg/kv/kvserver/kvflowcontrol/kvflowcontrol.go b/pkg/kv/kvserver/kvflowcontrol/kvflowcontrol.go index 94b5aef30a21..09bb9a44cf7e 100644 --- a/pkg/kv/kvserver/kvflowcontrol/kvflowcontrol.go +++ b/pkg/kv/kvserver/kvflowcontrol/kvflowcontrol.go @@ -19,8 +19,8 @@ import ( "github.com/cockroachdb/cockroach/pkg/settings/cluster" "github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb" "github.com/cockroachdb/cockroach/pkg/util/ctxutil" + "github.com/cockroachdb/cockroach/pkg/util/humanizeutil" "github.com/cockroachdb/redact" - "github.com/dustin/go-humanize" ) // Enabled determines whether we use flow control for replication traffic in KV. @@ -70,12 +70,12 @@ func (m ModeT) String() string { } // SafeFormat implements the redact.SafeFormatter interface. -func (m ModeT) SafeFormat(p redact.SafePrinter, verb rune) { +func (m ModeT) SafeFormat(p redact.SafePrinter, _ rune) { if s, ok := modeDict[m]; ok { - p.Print(s) + p.SafeString(redact.SafeString(s)) return } - p.Print("unknown-mode") + p.SafeString("unknown-mode") } // RegularTokensPerStream determines the flow tokens available for regular work @@ -415,13 +415,12 @@ func (t Tokens) String() string { } // SafeFormat implements the redact.SafeFormatter interface. -func (t Tokens) SafeFormat(p redact.SafePrinter, verb rune) { - sign := "+" +func (t Tokens) SafeFormat(p redact.SafePrinter, _ rune) { if t < 0 { - sign = "-" - t = -t + p.SafeString(humanizeutil.IBytes(int64(t))) + return } - p.Printf("%s%s", sign, humanize.IBytes(uint64(t))) + p.Printf("+%s", humanizeutil.IBytes(int64(t))) } func (s Stream) String() string { @@ -429,12 +428,12 @@ func (s Stream) String() string { } // SafeFormat implements the redact.SafeFormatter interface. -func (s Stream) SafeFormat(p redact.SafePrinter, verb rune) { +func (s Stream) SafeFormat(p redact.SafePrinter, _ rune) { tenantSt := s.TenantID.String() if s.TenantID.IsSystem() { tenantSt = "1" } - p.Printf("t%s/s%s", tenantSt, s.StoreID.String()) + p.Printf("t%s/s%s", redact.SafeString(tenantSt), s.StoreID) } var raftAdmissionMetaKey = ctxutil.RegisterFastValueKey() diff --git a/pkg/kv/kvserver/kvflowcontrol/replica_rac2/processor.go b/pkg/kv/kvserver/kvflowcontrol/replica_rac2/processor.go index 7332309967e9..f5798189eb67 100644 --- a/pkg/kv/kvserver/kvflowcontrol/replica_rac2/processor.go +++ b/pkg/kv/kvserver/kvflowcontrol/replica_rac2/processor.go @@ -922,7 +922,7 @@ func (p *processorImpl) AdmitRaftEntriesRaftMuLocked(ctx context.Context, e rac2 if isV2Encoding { log.Infof(ctx, "decoded v2 raft admission meta below-raft: pri=%v create-time=%d "+ - "proposer=n%v receiver=[n%d,s%v] tenant=t%d tokens≈%d "+ + "proposer=n%v receiver=[n%d,s%v] tenant=t%d tokens≈%v "+ "sideloaded=%t raft-entry=%d/%d lead-v2=%v", raftpb.Priority(meta.AdmissionPriority), meta.AdmissionCreateTime, @@ -939,7 +939,7 @@ func (p *processorImpl) AdmitRaftEntriesRaftMuLocked(ctx context.Context, e rac2 } else { log.Infof(ctx, "decoded v1 raft admission meta below-raft: pri=%v create-time=%d "+ - "proposer=n%v receiver=[n%d,s%v] tenant=t%d tokens≈%d "+ + "proposer=n%v receiver=[n%d,s%v] tenant=t%d tokens≈%v "+ "sideloaded=%t raft-entry=%d/%d lead-v2=%v", admissionpb.WorkPriority(meta.AdmissionPriority), meta.AdmissionCreateTime, diff --git a/pkg/sql/pgwire/auth_methods.go b/pkg/sql/pgwire/auth_methods.go index dc162c7bcf34..d9b3d9caecc1 100644 --- a/pkg/sql/pgwire/auth_methods.go +++ b/pkg/sql/pgwire/auth_methods.go @@ -350,7 +350,7 @@ func scramAuthenticator( // SASLResponse messages contain just the SASL payload. // rb := pgwirebase.ReadBuffer{Msg: resp} - reqMethod, err := rb.GetString() + reqMethod, err := rb.GetUnsafeString() if err != nil { c.LogAuthFailed(ctx, eventpb.AuthFailReason_PRE_HOOK_ERROR, err) return err diff --git a/pkg/sql/pgwire/conn.go b/pkg/sql/pgwire/conn.go index 91d80e64d430..738bee3b894d 100644 --- a/pkg/sql/pgwire/conn.go +++ b/pkg/sql/pgwire/conn.go @@ -346,7 +346,7 @@ func (c *conn) handleSimpleQuery( timeReceived crtime.Mono, unqualifiedIntSize *types.T, ) error { - query, err := buf.GetString() + query, err := buf.GetUnsafeString() if err != nil { return c.stmtBuf.Push(ctx, sql.SendError{Err: err}) } @@ -493,11 +493,11 @@ func (c *conn) handleSimpleQuery( // the connection should be considered toast. func (c *conn) handleParse(ctx context.Context, nakedIntSize *types.T) error { telemetry.Inc(sqltelemetry.ParseRequestCounter) - name, err := c.readBuf.GetString() + name, err := c.readBuf.GetUnsafeString() if err != nil { return c.stmtBuf.Push(ctx, sql.SendError{Err: err}) } - query, err := c.readBuf.GetString() + query, err := c.readBuf.GetUnsafeString() if err != nil { return c.stmtBuf.Push(ctx, sql.SendError{Err: err}) } @@ -608,7 +608,7 @@ func (c *conn) handleDescribe(ctx context.Context) error { if err != nil { return c.stmtBuf.Push(ctx, sql.SendError{Err: err}) } - name, err := c.readBuf.GetString() + name, err := c.readBuf.GetUnsafeString() if err != nil { return c.stmtBuf.Push(ctx, sql.SendError{Err: err}) } @@ -628,7 +628,7 @@ func (c *conn) handleClose(ctx context.Context) error { if err != nil { return c.stmtBuf.Push(ctx, sql.SendError{Err: err}) } - name, err := c.readBuf.GetString() + name, err := c.readBuf.GetUnsafeString() if err != nil { return c.stmtBuf.Push(ctx, sql.SendError{Err: err}) } @@ -650,11 +650,11 @@ var formatCodesAllText = []pgwirebase.FormatCode{pgwirebase.FormatText} // the connection should be considered toast. func (c *conn) handleBind(ctx context.Context) error { telemetry.Inc(sqltelemetry.BindRequestCounter) - portalName, err := c.readBuf.GetString() + portalName, err := c.readBuf.GetUnsafeString() if err != nil { return c.stmtBuf.Push(ctx, sql.SendError{Err: err}) } - statementName, err := c.readBuf.GetString() + statementName, err := c.readBuf.GetUnsafeString() if err != nil { return c.stmtBuf.Push(ctx, sql.SendError{Err: err}) } @@ -776,7 +776,7 @@ func (c *conn) handleExecute( ctx context.Context, timeReceived crtime.Mono, followedBySync bool, ) error { telemetry.Inc(sqltelemetry.ExecuteRequestCounter) - portalName, err := c.readBuf.GetString() + portalName, err := c.readBuf.GetUnsafeString() if err != nil { return c.stmtBuf.Push(ctx, sql.SendError{Err: err}) } diff --git a/pkg/sql/pgwire/pgwirebase/encoding.go b/pkg/sql/pgwire/pgwirebase/encoding.go index 740d675705b7..9297cf8aa5f6 100644 --- a/pkg/sql/pgwire/pgwirebase/encoding.go +++ b/pkg/sql/pgwire/pgwirebase/encoding.go @@ -213,8 +213,10 @@ func (b *ReadBuffer) ReadTypedMsg(rd BufferedReader) (ClientMessageType, int, er return ClientMessageType(typ), n, err } -// GetString reads a null-terminated string. -func (b *ReadBuffer) GetString() (string, error) { +// GetUnsafeString reads a null-terminated string as a reference. +// Note: The underlying buffer will be prevented from GCing, so long lived +// objects should never use this. +func (b *ReadBuffer) GetUnsafeString() (string, error) { pos := bytes.IndexByte(b.Msg, 0) if pos == -1 { return "", NewProtocolViolationErrorf("NUL terminator not found") @@ -226,6 +228,16 @@ func (b *ReadBuffer) GetString() (string, error) { return s, nil } +// GetSafeString reads a null-terminated string as a copy of the original data +// out. +func (b *ReadBuffer) GetSafeString() (string, error) { + s, err := b.GetUnsafeString() + if err != nil { + return "", err + } + return strings.Clone(s), nil +} + // GetPrepareType returns the buffer's contents as a PrepareType. func (b *ReadBuffer) GetPrepareType() (PrepareType, error) { v, err := b.GetBytes(1) diff --git a/pkg/sql/pgwire/pre_serve_options.go b/pkg/sql/pgwire/pre_serve_options.go index c8559618ebe4..d5e7a61c328d 100644 --- a/pkg/sql/pgwire/pre_serve_options.go +++ b/pkg/sql/pgwire/pre_serve_options.go @@ -54,7 +54,9 @@ func parseClientProvidedSessionParameters( hasTenantSelectOption := false for { // Read a key-value pair from the client. - key, err := buf.GetString() + // Note: GetSafeString is used since the key/value will live well past the + // life of the message. + key, err := buf.GetSafeString() if err != nil { return args, pgerror.Wrap( err, pgcode.ProtocolViolation, @@ -65,7 +67,7 @@ func parseClientProvidedSessionParameters( // End of parameter list. break } - value, err := buf.GetString() + value, err := buf.GetSafeString() if err != nil { return args, pgerror.Wrapf( err, pgcode.ProtocolViolation, @@ -75,10 +77,6 @@ func parseClientProvidedSessionParameters( // Case-fold for the key for easier comparison. key = strings.ToLower(key) - // Intentionally clone the string from above, so that the ReaderBuffer life - // is limited. Otherwise, the buffer will remain allocated for the life of - // the connection. - value = strings.Clone(value) // Load the parameter. switch key { diff --git a/pkg/util/metamorphic/BUILD.bazel b/pkg/util/metamorphic/BUILD.bazel index f9fd4b5416e8..4908074ba379 100644 --- a/pkg/util/metamorphic/BUILD.bazel +++ b/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", + ], ) diff --git a/pkg/util/metamorphic/constants.go b/pkg/util/metamorphic/constants.go index 78a43416b932..00b3914203de 100644 --- a/pkg/util/metamorphic/constants.go +++ b/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] + } } diff --git a/pkg/util/metamorphic/constants_test.go b/pkg/util/metamorphic/constants_test.go index 6c075ab57ede..b2d3ce20d0a7 100644 --- a/pkg/util/metamorphic/constants_test.go +++ b/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) +} diff --git a/pkg/util/metamorphic/main_test.go b/pkg/util/metamorphic/main_test.go new file mode 100644 index 000000000000..98fce9e7e151 --- /dev/null +++ b/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()) +}