util/rands: add a cheap non-escaping rand typeraggi/rand

math/rand and the current v2 proposal retain a source interface that
means they always heap allocate until some future compiler changes this
constraint. This type is one that can be used on-stack for use cases
like a cheap in-place shuffle with an on-stack random generator.

Updates #17243
Signed-off-by: James Tucker <james@tailscale.com>

3 files changed, 263 insertions, 0 deletions

diff --git a/cmd/tailscaled/depaware.txt b/cmd/tailscaled/depaware.txt
index 52387e849..a55c0e7aa 100644
--- a/cmd/tailscaled/depaware.txt
+++ b/cmd/tailscaled/depaware.txt
@@ -416,6 +416,7 @@ tailscale.com/cmd/tailscaled dependencies: (generated by github.com/tailscale/de
   LD    golang.org/x/crypto/ssh                                      from tailscale.com/ssh/tailssh+
         golang.org/x/exp/constraints                                 from github.com/dblohm7/wingoes/pe+
         golang.org/x/exp/maps                                        from tailscale.com/wgengine/magicsock+
+        golang.org/x/exp/rand                                        from tailscale.com/util/rands
         golang.org/x/net/bpf                                         from github.com/mdlayher/genetlink+
         golang.org/x/net/dns/dnsmessage                              from net+
         golang.org/x/net/http/httpguts                               from golang.org/x/net/http2+
diff --git a/util/rands/cheap.go b/util/rands/cheap.go
new file mode 100644
index 000000000..a80942f67
--- /dev/null
+++ b/util/rands/cheap.go
@@ -0,0 +1,66 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package rands
+
+import (
+	exprand "golang.org/x/exp/rand"
+)
+
+// A Rand is a source of random numbers. It is extremely cheap to create and
+// seed on the stack, and always uses the PCG random number generator.
+type Rand struct {
+	src exprand.PCGSource
+}
+
+// NewRand returns a new Rand with the given seed.
+func NewRand(seed uint64) Rand {
+	var r Rand
+	r.Seed(seed)
+	return r
+}
+
+// Seed uses the provided seed value to reinitialize the generator to a
+// deterministic state.
+// Seed should not be called concurrently with any other Rand method.
+func (r *Rand) Seed(seed uint64) {
+	r.src.Seed(seed)
+}
+
+// Uint64 returns a pseudo-random 64-bit integer as a uint64.
+func (r *Rand) Uint64() uint64 { return r.src.Uint64() }
+
+const maxUint64 = (1 << 64) - 1
+
+// Uint64n returns, as a uint64, a pseudo-random number in [0,n).
+// It is guaranteed more uniform than taking a Source value mod n
+// for any n that is not a power of 2.
+func (r *Rand) Uint64n(n uint64) uint64 {
+	if n&(n-1) == 0 { // n is power of two, can mask
+		if n == 0 {
+			panic("invalid argument to Uint64n")
+		}
+		return r.Uint64() & (n - 1)
+	}
+	// If n does not divide v, to avoid bias we must not use
+	// a v that is within maxUint64%n of the top of the range.
+	v := r.Uint64()
+	if v > maxUint64-n { // Fast check.
+		ceiling := maxUint64 - maxUint64%n
+		for v >= ceiling {
+			v = r.Uint64()
+		}
+	}
+
+	return v % n
+}
+
+// Intn returns, as an int, a non-negative pseudo-random number in [0,n).
+// It panics if n <= 0.
+func (r *Rand) Intn(n int) int {
+	if n <= 0 {
+		panic("invalid argument to Intn")
+	}
+	// TODO: Avoid some 64-bit ops to make it more efficient on 32-bit machines.
+	return int(r.Uint64n(uint64(n)))
+}
diff --git a/util/rands/cheap_test.go b/util/rands/cheap_test.go
new file mode 100644
index 000000000..6120d3a10
--- /dev/null
+++ b/util/rands/cheap_test.go
@@ -0,0 +1,196 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package rands
+
+import (
+	"math/rand"
+	"sync"
+	"testing"
+
+	exprand "golang.org/x/exp/rand"
+)
+
+var (
+	seed          uint64 = 8729831
+	numDraw              = 100
+	numGoroutines        = 5000
+)
+
+type workerPool struct {
+	job chan func()
+	res chan struct{}
+	wg  sync.WaitGroup
+}
+
+func (p *workerPool) Close() {
+	close(p.job)
+	p.wg.Wait()
+}
+
+func newWorkerPool() *workerPool {
+	pool := workerPool{
+		job: make(chan func(), 2<<20),
+		res: make(chan struct{}, 2<<20),
+		wg:  sync.WaitGroup{},
+	}
+	for i := 0; i < numGoroutines; i++ {
+		pool.wg.Add(1)
+		go func() {
+			defer pool.wg.Done()
+			for f := range pool.job {
+				f()
+				pool.res <- struct{}{}
+			}
+		}()
+	}
+	return &pool
+}
+
+var stdPool = sync.Pool{
+	New: func() any {
+		return rand.New(rand.NewSource(int64(seed)))
+	},
+}
+
+var expPool = sync.Pool{
+	New: func() any {
+		return exprand.New(exprand.NewSource(seed))
+	},
+}
+
+func BenchmarkStd(b *testing.B) {
+	pool := newWorkerPool()
+	defer pool.Close()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		pool.job <- func() {
+			rand.Seed(int64(seed))
+			for i := 0; i < numDraw; i++ {
+				rand.Intn(100)
+			}
+		}
+	}
+	for i := 0; i < b.N; i++ {
+		<-pool.res
+	}
+}
+
+func BenchmarkPCG(b *testing.B) {
+	pool := newWorkerPool()
+	defer pool.Close()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		pool.job <- func() {
+			exprand.Seed(seed)
+			for i := 0; i < numDraw; i++ {
+				exprand.Intn(100)
+			}
+		}
+	}
+	for i := 0; i < b.N; i++ {
+		<-pool.res
+	}
+}
+
+func BenchmarkStdPool(b *testing.B) {
+	pool := newWorkerPool()
+	defer pool.Close()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		pool.job <- func() {
+			r := stdPool.Get().(*rand.Rand)
+			defer stdPool.Put(r)
+
+			r.Seed(int64(seed))
+			for i := 0; i < numDraw; i++ {
+				r.Intn(100)
+			}
+		}
+	}
+	for i := 0; i < b.N; i++ {
+		<-pool.res
+	}
+}
+
+func BenchmarkPCGPool(b *testing.B) {
+	pool := newWorkerPool()
+	defer pool.Close()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		pool.job <- func() {
+			r := expPool.Get().(*exprand.Rand)
+			defer expPool.Put(r)
+
+			r.Seed(seed)
+			for i := 0; i < numDraw; i++ {
+				r.Intn(100)
+			}
+		}
+	}
+	for i := 0; i < b.N; i++ {
+		<-pool.res
+	}
+}
+
+func BenchmarkLocalStd(b *testing.B) {
+	pool := newWorkerPool()
+	defer pool.Close()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		pool.job <- func() {
+			r := rand.New(rand.NewSource(int64(seed)))
+			for i := 0; i < numDraw; i++ {
+				r.Intn(100)
+			}
+		}
+	}
+	for i := 0; i < b.N; i++ {
+		<-pool.res
+	}
+}
+
+func BenchmarkLocalPCG(b *testing.B) {
+	pool := newWorkerPool()
+	defer pool.Close()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		pool.job <- func() {
+			r := exprand.New(exprand.NewSource(seed))
+			for i := 0; i < numDraw; i++ {
+				r.Intn(100)
+			}
+		}
+	}
+	for i := 0; i < b.N; i++ {
+		<-pool.res
+	}
+}
+
+func BenchmarkStackRand(b *testing.B) {
+	pool := newWorkerPool()
+	defer pool.Close()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		pool.job <- func() {
+			r := NewRand(seed)
+			for i := 0; i < numDraw; i++ {
+				r.Intn(100)
+			}
+		}
+	}
+	for i := 0; i < b.N; i++ {
+		<-pool.res
+	}
+}
+
+func TestStackRandNoAllocs(t *testing.T) {
+	seed := rand.Uint64()
+	if n := testing.AllocsPerRun(1000, func() {
+		r := NewRand(seed)
+		_ = r.Intn(100)
+	}); n > 0 {
+		t.Errorf("Rand got %v allocs per run", n)
+	}
+
+}