Merge #149053 #149197

149053: sql: fix cancel leak in happy case in fingerprint_span r=yuzefovich a=yuzefovich

The recent fix introduced a possible leak of the cancel function (namely, in the happy case when the coordinator doesn't hit an early error, we would never call the cancellation function). We should just defer it unconditionally and might call it earlier on a shutdown due to an error.

Epic: None
Release note: None

149197: kvnemesis: log seed and rng access count r=miraradeva a=stevendanna

This new logging

1) Ensures the seed is captured as part of the kvnemesis output even if the test runner logs are somehow lost.

2) Allows us to spot when a seed might not be reproducing a problem because of a change in the number of random choices in KVNemesis.

Epic: none
Release note: None

Co-authored-by: Yahor Yuzefovich <[email protected]>
Co-authored-by: Steven Danna <[email protected]>

Author: 3 people

pkg/kv/kvnemesis/generator.go

@@ -14,6 +14,7 @@ import (
 	"math/rand"
 	"slices"
 	"sort"
+	"sync/atomic"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvnemesis/kvnemesisutil"
@@ -36,6 +37,9 @@ type GeneratorConfig struct {
 	NumNodes, NumReplicas int
 
 	BufferedWritesProb float64
+
+	SeedForLogging              int64
+	RandSourceCounterForLogging counter
 }
 
 // OperationConfig configures the relative probabilities of producing various
@@ -1993,3 +1997,40 @@ func releaseSavepoint(id int) Operation {
 func rollbackSavepoint(id int) Operation {
 	return Operation{SavepointRollback: &SavepointRollbackOperation{ID: int32(id)}}
 }
+
+type countingRandSource struct {
+	count atomic.Uint64
+	inner rand.Source64
+}
+
+type counter interface {
+	Count() uint64
+}
+
+// newCountingSource creates random source that counts how many times it was
+// called for logging purposes.
+func newCountingSource(inner rand.Source64) *countingRandSource {
+	return &countingRandSource{
+		inner: inner,
+	}
+}
+
+func (c *countingRandSource) Count() uint64 {
+	return c.count.Load()
+}
+
+func (c *countingRandSource) Int63() int64 {
+	c.count.Add(1)
+	return c.inner.Int63()
+}
+
+func (c *countingRandSource) Uint64() uint64 {
+	c.count.Add(1)
+	return c.inner.Uint64()
+}
+
+func (c *countingRandSource) Seed(seed int64) {
+	// We assume that seed invalidates the count.
+	c.count.Store(0)
+	c.inner.Seed(seed)
+}

pkg/kv/kvnemesis/kvnemesis.go

@@ -177,7 +177,13 @@ func RunNemesis(
 			failuresFile = l(ctx, "failures", "%s", &buf)
 		}
 
-		reproFile := l(ctx, "repro.go", "// Reproduction steps:\n%s", printRepro(stepsByWorker))
+		reproFile := l(ctx, "repro.go", `// Seed: %d
+// Calls to Random Source: %d
+// Reproduction steps:
+%s`,
+			config.SeedForLogging,
+			config.RandSourceCounterForLogging.Count(),
+			printRepro(stepsByWorker))
 		rangefeedFile := l(ctx, "kvs-rangefeed.txt", "kvs (recorded from rangefeed):\n%s", kvs.DebugPrint("  "))
 		kvsFile := "<error>"
 		var scanKVs []kv.KeyValue

pkg/kv/kvnemesis/kvnemesis_test.go

@@ -210,16 +210,17 @@ func randWithSeed(
 		Logf(string, ...interface{})
 		Helper()
 	}, seedOrZero int64,
-) *rand.Rand {
+) (*rand.Rand, counter, int64) {
 	t.Helper()
-	var rng *rand.Rand
+	var rngSource rand.Source
 	if seedOrZero > 0 {
-		rng = rand.New(rand.NewSource(seedOrZero))
+		rngSource = rand.NewSource(seedOrZero)
 	} else {
-		rng, seedOrZero = randutil.NewTestRand()
+		rngSource, seedOrZero = randutil.NewTestRandSource()
 	}
 	t.Logf("seed: %d", seedOrZero)
-	return rng
+	countingSource := newCountingSource(rngSource.(rand.Source64))
+	return rand.New(countingSource), countingSource, seedOrZero
 }
 
 type ti interface {
@@ -445,7 +446,7 @@ func testKVNemesisImpl(t *testing.T, cfg kvnemesisTestCfg) {
 
 	// Can set a seed here for determinism. This works best when the seed was
 	// obtained with cfg.concurrency=1.
-	rng := randWithSeed(t, cfg.seedOverride)
+	rng, countingSource, seed := randWithSeed(t, cfg.seedOverride)
 
 	// 4 nodes so we have somewhere to move 3x replicated ranges to.
 	ctx := context.Background()
@@ -466,6 +467,10 @@ func testKVNemesisImpl(t *testing.T, cfg kvnemesisTestCfg) {
 	config.NumNodes = cfg.numNodes
 	config.NumReplicas = 3
 	config.BufferedWritesProb = cfg.bufferedWriteProb
+
+	config.SeedForLogging = seed
+	config.RandSourceCounterForLogging = countingSource
+
 	if config.NumReplicas > cfg.numNodes {
 		config.NumReplicas = cfg.numNodes
 	}

pkg/sql/fingerprint_span.go

@@ -139,6 +139,7 @@ func (p *planner) fingerprintSpanFanout(
 			// the coordinator encounters an error.
 			workCh := make(chan roachpb.Span)
 			ctx, cancel := context.WithCancel(ctx)
+			defer cancel()
 
 			grp := ctxgroup.WithContext(ctx)
 			for range maxWorkerCount {

pkg/util/randutil/rand.go

@@ -115,16 +115,27 @@ func NewPseudoRandWithGlobalSeed() (*rand.Rand, int64) {
 // seed. This rand.Rand is useful in testing to produce deterministic,
 // reproducible behavior.
 func NewTestRand() (*rand.Rand, int64) {
-	return newTestRandImpl(rand.NewSource)
+	src, seed := newTestRandSourceImpl(rand.NewSource)
+	return rand.New(src), seed
+}
+
+// NewTestRandSource returns a math/rand.Source64 seeded from rng, which is seeded
+// with the global seed. If the caller is a test with a different path-qualified
+// name than the previous caller, rng is reseeded from the global seed. This
+// random source is useful in testing to produce deterministic, reproducible
+// behavior.
+func NewTestRandSource() (rand.Source, int64) {
+	return newTestRandSourceImpl(rand.NewSource)
 }
 
 // NewLockedTestRand is identical to NewTestRand but returned rand.Rand is using
 // thread safe underlying source.
 func NewLockedTestRand() (*rand.Rand, int64) {
-	return newTestRandImpl(NewLockedSource)
+	src, seed := newTestRandSourceImpl(NewLockedSource)
+	return rand.New(src), seed
 }
 
-func newTestRandImpl(f func(int64) rand.Source) (*rand.Rand, int64) {
+func newTestRandSourceImpl(f func(int64) rand.Source) (rand.Source, int64) {
 	mtx.Lock()
 	defer mtx.Unlock()
 	fxn := getTestName()
@@ -136,7 +147,7 @@ func newTestRandImpl(f func(int64) rand.Source) (*rand.Rand, int64) {
 		rng = rand.New(f(globalSeed))
 	}
 	seed := rng.Int63()
-	return rand.New(f(seed)), seed
+	return f(seed), seed
 }
 
 // NewTestRandWithSeed returns an instance of math/rand.Rand, similar to