From d0f3fa7d7eb63b604a57936e44e39d0d868aa69a Mon Sep 17 00:00:00 2001
From: James Tucker <james@tailscale.com>
Date: Tue, 9 Apr 2024 13:02:18 -0700
Subject: [PATCH] util/fastuuid: add a more efficient uuid generator

This still generates github.com/google/uuid UUID objects, but does so
using a ChaCha8 CSPRNG from the stdlib rand/v2 package. The public API
is backed by a sync.Pool to provide good performance in highly
concurrent operation.

Under high load the read API produces a lot of extra garbage and
overhead by way of temporaries and syscalls. This implementation reduces
both to minimal levels, and avoids any long held global lock by
utilizing sync.Pool.

Updates tailscale/corp#18266
Updates tailscale/corp#19054

Signed-off-by: James Tucker <james@tailscale.com>
---
 util/fastuuid/fastuuid.go      | 56 ++++++++++++++++++++++++++
 util/fastuuid/fastuuid_test.go | 72 ++++++++++++++++++++++++++++++++++
 2 files changed, 128 insertions(+)
 create mode 100644 util/fastuuid/fastuuid.go
 create mode 100644 util/fastuuid/fastuuid_test.go

diff --git a/util/fastuuid/fastuuid.go b/util/fastuuid/fastuuid.go
new file mode 100644
index 000000000..4b115ea4e
--- /dev/null
+++ b/util/fastuuid/fastuuid.go
@@ -0,0 +1,56 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+// Package fastuuid implements a UUID construction using an in process CSPRNG.
+package fastuuid
+
+import (
+	crand "crypto/rand"
+	"encoding/binary"
+	"io"
+	"math/rand/v2"
+	"sync"
+
+	"github.com/google/uuid"
+)
+
+// NewUUID returns a new UUID using a pool of generators, good for highly
+// concurrent use.
+func NewUUID() uuid.UUID {
+	g := pool.Get().(*generator)
+	defer pool.Put(g)
+	return g.newUUID()
+}
+
+var pool = sync.Pool{
+	New: func() any {
+		return newGenerator()
+	},
+}
+
+type generator struct {
+	rng rand.ChaCha8
+}
+
+func seed() [32]byte {
+	var r [32]byte
+	if _, err := io.ReadFull(crand.Reader, r[:]); err != nil {
+		panic(err)
+	}
+	return r
+}
+
+func newGenerator() *generator {
+	return &generator{
+		rng: *rand.NewChaCha8(seed()),
+	}
+}
+
+func (g *generator) newUUID() uuid.UUID {
+	var u uuid.UUID
+	binary.NativeEndian.PutUint64(u[:8], g.rng.Uint64())
+	binary.NativeEndian.PutUint64(u[8:], g.rng.Uint64())
+	u[6] = (u[6] & 0x0f) | 0x40 // Version 4
+	u[8] = (u[8] & 0x3f) | 0x80 // Variant 10
+	return u
+}
diff --git a/util/fastuuid/fastuuid_test.go b/util/fastuuid/fastuuid_test.go
new file mode 100644
index 000000000..b514d6daa
--- /dev/null
+++ b/util/fastuuid/fastuuid_test.go
@@ -0,0 +1,72 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package fastuuid
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+)
+
+func TestNewUUID(t *testing.T) {
+	g := pool.Get().(*generator)
+	defer pool.Put(g)
+	u := g.newUUID()
+	if u[6] != (u[6]&0x0f)|0x40 {
+		t.Errorf("version bits are incorrect")
+	}
+	if u[8] != (u[8]&0x3f)|0x80 {
+		t.Errorf("variant bits are incorrect")
+	}
+}
+
+func BenchmarkBasic(b *testing.B) {
+	b.Run("NewUUID", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			NewUUID()
+		}
+	})
+
+	b.Run("uuid.New-unpooled", func(b *testing.B) {
+		uuid.DisableRandPool()
+		for i := 0; i < b.N; i++ {
+			uuid.New()
+		}
+	})
+
+	b.Run("uuid.New-pooled", func(b *testing.B) {
+		uuid.EnableRandPool()
+		for i := 0; i < b.N; i++ {
+			uuid.New()
+		}
+	})
+}
+
+func BenchmarkParallel(b *testing.B) {
+	b.Run("NewUUID", func(b *testing.B) {
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				NewUUID()
+			}
+		})
+	})
+
+	b.Run("uuid.New-unpooled", func(b *testing.B) {
+		uuid.DisableRandPool()
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				uuid.New()
+			}
+		})
+	})
+
+	b.Run("uuid.New-pooled", func(b *testing.B) {
+		uuid.EnableRandPool()
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				uuid.New()
+			}
+		})
+	})
+}