workload: reduce RNG allocations

PR cockroachdb#143626 removed `golang.org/x/exp/rand` in favor of `math/rand/v2`.
The number of heap allocations grew dramatically due to new calls to
`rand.New` and `rand.NewPCG`. This commit eliminates these allocations
by reusing allocations of `rand.Rand` and `rand.PCG`.

```
name                           old time/op    new time/op    delta
InitialData/tpcc/warehouses=1     103ms ± 0%      91ms ± 0%  -11.96%  (p=0.016 n=4+5)

name                           old speed      new speed      delta
InitialData/tpcc/warehouses=1  2.13GB/s ± 0%  2.50GB/s ± 0%  +17.58%  (p=0.016 n=4+5)

name                           old alloc/op   new alloc/op   delta
InitialData/tpcc/warehouses=1    10.3MB ± 0%     0.1MB ± 0%  -99.24%  (p=0.008 n=5+5)

name                           old allocs/op  new allocs/op  delta
InitialData/tpcc/warehouses=1      640k ± 0%        0k ± 0%  -99.97%  (p=0.008 n=5+5)
```

Release note: None

Author: mgartner

pkg/workload/tpcc/generate.go

@@ -72,7 +72,7 @@ var itemTypes = []*types.T{
 func (w *tpcc) tpccItemInitialRowBatch(rowIdx int, cb coldata.Batch, a *bufalloc.ByteAllocator) {
 	l := w.localsPool.Get().(*generateLocals)
 	defer w.localsPool.Put(l)
-	l.rng.Rand = rand.New(rand.NewPCG(RandomSeed.Seed(), uint64(rowIdx)))
+	l.rng.Seed(RandomSeed.Seed(), uint64(rowIdx))
 	ao := aCharsOffset(l.rng.IntN(len(aCharsAlphabet)))
 
 	iID := rowIdx + 1
@@ -115,7 +115,7 @@ func (w *tpcc) tpccWarehouseInitialRowBatch(
 ) {
 	l := w.localsPool.Get().(*generateLocals)
 	defer w.localsPool.Put(l)
-	l.rng.Rand = rand.New(rand.NewPCG(RandomSeed.Seed(), uint64(rowIdx)))
+	l.rng.Seed(RandomSeed.Seed(), uint64(rowIdx))
 	no := numbersOffset(l.rng.IntN(len(numbersAlphabet)))
 	lo := lettersOffset(l.rng.IntN(len(lettersAlphabet)))
 
@@ -173,7 +173,7 @@ var stockTypes = []*types.T{
 func (w *tpcc) tpccStockInitialRowBatch(rowIdx int, cb coldata.Batch, a *bufalloc.ByteAllocator) {
 	l := w.localsPool.Get().(*generateLocals)
 	defer w.localsPool.Put(l)
-	l.rng.Rand = rand.New(rand.NewPCG(RandomSeed.Seed(), uint64(rowIdx)))
+	l.rng.Seed(RandomSeed.Seed(), uint64(rowIdx))
 	ao := aCharsOffset(l.rng.IntN(len(aCharsAlphabet)))
 
 	sID := (rowIdx % numStockPerWarehouse) + 1
@@ -245,7 +245,7 @@ func (w *tpcc) tpccDistrictInitialRowBatch(
 ) {
 	l := w.localsPool.Get().(*generateLocals)
 	defer w.localsPool.Put(l)
-	l.rng.Rand = rand.New(rand.NewPCG(RandomSeed.Seed(), uint64(rowIdx)))
+	l.rng.Seed(RandomSeed.Seed(), uint64(rowIdx))
 	ao := aCharsOffset(l.rng.IntN(len(aCharsAlphabet)))
 	no := numbersOffset(l.rng.IntN(len(numbersAlphabet)))
 	lo := lettersOffset(l.rng.IntN(len(lettersAlphabet)))
@@ -318,7 +318,7 @@ func (w *tpcc) tpccCustomerInitialRowBatch(
 ) {
 	l := w.localsPool.Get().(*generateLocals)
 	defer w.localsPool.Put(l)
-	l.rng.Rand = rand.New(rand.NewPCG(RandomSeed.Seed(), uint64(rowIdx)))
+	l.rng.Seed(RandomSeed.Seed(), uint64(rowIdx))
 	ao := aCharsOffset(l.rng.IntN(len(aCharsAlphabet)))
 	no := numbersOffset(l.rng.IntN(len(numbersAlphabet)))
 	lo := lettersOffset(l.rng.IntN(len(lettersAlphabet)))
@@ -416,7 +416,7 @@ var historyTypes = []*types.T{
 func (w *tpcc) tpccHistoryInitialRowBatch(rowIdx int, cb coldata.Batch, a *bufalloc.ByteAllocator) {
 	l := w.localsPool.Get().(*generateLocals)
 	defer w.localsPool.Put(l)
-	l.rng.Rand = rand.New(rand.NewPCG(RandomSeed.Seed(), uint64(rowIdx)))
+	l.rng.Seed(RandomSeed.Seed(), uint64(rowIdx))
 	ao := aCharsOffset(l.rng.IntN(len(aCharsAlphabet)))
 
 	// This used to be a V4 uuid made through the normal `uuid.MakeV4`
@@ -481,7 +481,7 @@ func (w *tpcc) tpccOrderInitialRowBatch(rowIdx int, cb coldata.Batch, a *bufallo
 
 	// NB: numOrderLines is not allowed to use precomputed random data, make sure
 	// it stays that way. See 4.3.2.1.
-	l.rng.Rand = rand.New(rand.NewPCG(RandomSeed.Seed(), uint64(rowIdx)))
+	l.rng.Seed(RandomSeed.Seed(), uint64(rowIdx))
 	numOrderLines := randInt(l.rng.Rand, minOrderLinesPerOrder, maxOrderLinesPerOrder)
 
 	oID := (rowIdx % numOrdersPerDistrict) + 1
@@ -602,7 +602,7 @@ func (w *tpcc) tpccOrderLineInitialRowBatch(
 
 	// NB: numOrderLines is not allowed to use precomputed random data, make sure
 	// it stays that way. See 4.3.2.1.
-	l.rng.Rand = rand.New(rand.NewPCG(RandomSeed.Seed(), uint64(orderRowIdx)))
+	l.rng.Seed(RandomSeed.Seed(), uint64(orderRowIdx))
 	numOrderLines := int(randInt(l.rng.Rand, minOrderLinesPerOrder, maxOrderLinesPerOrder))
 
 	// NB: There is one batch of order_line rows per order

pkg/workload/tpcc/random.go

@@ -32,6 +32,7 @@ const numbersAlphabet = `1234567890`
 
 type tpccRand struct {
 	*rand.Rand
+	*rand.PCG
 
 	aChars, letters, numbers workloadimpl.PrecomputedRand
 }
@@ -40,6 +41,18 @@ type aCharsOffset int
 type lettersOffset int
 type numbersOffset int
 
+// Seed resets the RNG with the given seeds.
+func (rng *tpccRand) Seed(seed1, seed2 uint64) {
+	if rng.PCG != nil {
+		// rand.Rand has no state other than the source, so we don't need to
+		// recreate it.
+		rng.PCG.Seed(seed1, seed2)
+		return
+	}
+	rng.PCG = rand.NewPCG(seed1, seed2)
+	rng.Rand = rand.New(rng.PCG)
+}
+
 func randStringFromAlphabet(
 	rng *rand.Rand,
 	a *bufalloc.ByteAllocator,