From 524ca6a3a3fbe89ef722eb270a969f951a2eb115 Mon Sep 17 00:00:00 2001
From: Paul Wells <paulwe@gmail.com>
Date: Sun, 21 Apr 2024 13:07:37 -0700
Subject: [PATCH] reduce allocations from guids (#692)

* reduce allocations from guids

* fix

* fix

* less lock thrashing
---
 utils/id.go      | 107 +++++++++++++++++++++++++++++++++++++++++------
 utils/id_test.go |  10 +++++
 2 files changed, 104 insertions(+), 13 deletions(-)

diff --git a/utils/id.go b/utils/id.go
index bcef4d05..ad08248c 100644
--- a/utils/id.go
+++ b/utils/id.go
@@ -15,14 +15,19 @@
 package utils
 
 import (
+	"crypto/rand"
 	"crypto/sha1"
+	"encoding/binary"
 	"fmt"
 	"os"
+	"sync"
+	"time"
 
 	"github.com/jxskiss/base62"
 	"github.com/lithammer/shortuuid/v4"
 
 	"github.com/livekit/protocol/livekit"
+	"github.com/livekit/protocol/logger"
 )
 
 const GuidSize = 12
@@ -45,8 +50,10 @@ const (
 	AgentJobPrefix        = "AJ_"
 )
 
+var defaultGuidGenerator = newGuidGenerator(4096, GuidSize+10)
+
 func NewGuid(prefix string) string {
-	return prefix + shortuuid.New()[:GuidSize]
+	return defaultGuidGenerator.NewGuid(prefix)
 }
 
 // HashedID creates a hashed ID from a unique string
@@ -66,17 +73,91 @@ func LocalNodeID() (string, error) {
 	return fmt.Sprintf("%s%s", NodePrefix, HashedID(hostname)[:8]), nil
 }
 
-var b62Index = newB62Index()
-var b62Chars = []byte(shortuuid.DefaultAlphabet)
+var b57Index = newB57Index()
+var b57Chars = []byte(shortuuid.DefaultAlphabet)
 
-func newB62Index() [256]byte {
+func newB57Index() [256]byte {
 	var index [256]byte
-	for i := 0; i < len(b62Chars); i++ {
-		index[b62Chars[i]] = byte(i)
+	for i := 0; i < len(b57Chars); i++ {
+		index[b57Chars[i]] = byte(i)
 	}
 	return index
 }
 
+type guidGenerator struct {
+	pool sync.Pool
+	mu   sync.Mutex
+	buf  []byte
+	pos  int
+}
+
+func newGuidGenerator(bufSize, scratchSize int) *guidGenerator {
+	return &guidGenerator{
+		pool: sync.Pool{
+			New: func() any {
+				b := make([]byte, scratchSize)
+				return &b
+			},
+		},
+		buf: make([]byte, bufSize),
+		pos: bufSize,
+	}
+}
+
+func (g *guidGenerator) refillBuf() {
+	for {
+		_, err := rand.Read(g.buf)
+		if err == nil {
+			return
+		}
+
+		logger.Errorw("unable to refill guid buffer", err)
+		time.Sleep(time.Millisecond)
+	}
+}
+
+func (g *guidGenerator) uint32() uint32 {
+	if g.pos == len(g.buf) {
+		g.refillBuf()
+		g.pos = 0
+	}
+
+	n := binary.BigEndian.Uint32(g.buf[g.pos:])
+	g.pos += 4
+
+	return n
+}
+
+func (g *guidGenerator) readIDChars(b []byte) {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	var n int
+	for {
+		r := g.uint32()
+		for i := 0; i < 5; i++ {
+			if int(r&0x3f) < len(b57Chars) {
+				b[n] = b57Chars[r&0x3f]
+				n++
+				if n == len(b) {
+					return
+				}
+			}
+			r >>= 6
+		}
+	}
+}
+
+func (g *guidGenerator) NewGuid(prefix string) string {
+	b := g.pool.Get().(*[]byte)
+	defer g.pool.Put(b)
+
+	*b = append((*b)[:0], make([]byte, len(prefix)+GuidSize)...)
+	copy(*b, prefix)
+	g.readIDChars((*b)[len(prefix):])
+	return string(*b)
+}
+
 func guidPrefix[T livekit.Guid]() string {
 	var id T
 	switch any(id).(type) {
@@ -97,9 +178,9 @@ func MarshalGuid[T livekit.Guid](id T) livekit.GuidBlock {
 	for i := 0; i < 3; i++ {
 		j := i * 3
 		k := i * 4
-		b[j] = b62Index[idb[k]]<<2 | b62Index[idb[k+1]]>>4
-		b[j+1] = b62Index[idb[k+1]]<<4 | b62Index[idb[k+2]]>>2
-		b[j+2] = b62Index[idb[k+2]]<<6 | b62Index[idb[k+3]]
+		b[j] = b57Index[idb[k]]<<2 | b57Index[idb[k+1]]>>4
+		b[j+1] = b57Index[idb[k+1]]<<4 | b57Index[idb[k+2]]>>2
+		b[j+2] = b57Index[idb[k+2]]<<6 | b57Index[idb[k+3]]
 	}
 	return b
 }
@@ -112,10 +193,10 @@ func UnmarshalGuid[T livekit.Guid](b livekit.GuidBlock) T {
 	for i := 0; i < 3; i++ {
 		j := i * 3
 		k := i * 4
-		idb[k] = b62Chars[b[j]>>2]
-		idb[k+1] = b62Chars[(b[j]&3)<<4|b[j+1]>>4]
-		idb[k+2] = b62Chars[(b[j+1]&15)<<2|b[j+2]>>6]
-		idb[k+3] = b62Chars[b[j+2]&63]
+		idb[k] = b57Chars[b[j]>>2]
+		idb[k+1] = b57Chars[(b[j]&3)<<4|b[j+1]>>4]
+		idb[k+2] = b57Chars[(b[j+1]&15)<<2|b[j+2]>>6]
+		idb[k+3] = b57Chars[b[j+2]&63]
 	}
 	return T(id)
 }
diff --git a/utils/id_test.go b/utils/id_test.go
index 2655acc3..9a15c3d4 100644
--- a/utils/id_test.go
+++ b/utils/id_test.go
@@ -30,3 +30,13 @@ func TestMarshalUnmarshalGuid(t *testing.T) {
 	require.EqualValues(t, id0, id1)
 	require.EqualValues(t, b0, b1)
 }
+
+func BenchmarkNewGuid(b *testing.B) {
+	b.Run("new", func(b *testing.B) {
+		var guid string
+		for i := 0; i < b.N; i++ {
+			guid = NewGuid(TrackPrefix)
+		}
+		_ = guid
+	})
+}