Rework bufpool to be a little more efficient, I hope

pdu/bufpool.go

@@ -1,48 +1,58 @@
 package pdu
 
 import (
-    "sync"
+	"sync"
 )
 
 type BufferPoolManager struct {
-    pools map[uint]*sync.Pool
-    mu    sync.Mutex
+	pools map[uint]*sync.Pool
+	mu    sync.RWMutex
 }
 
 func NewBufferPoolManager() *BufferPoolManager {
-    return &BufferPoolManager{
-        pools: make(map[uint]*sync.Pool),
-    }
+	return &BufferPoolManager{
+		pools: make(map[uint]*sync.Pool),
+	}
 }
 
-func (bpm *BufferPoolManager) GetBuffer(size uint) *([]uint8) {
-    bpm.mu.Lock()
-    pool, exists := bpm.pools[size]
-    if !exists {
-        pool = &sync.Pool{
-            New: func() interface{} {
-                buf := make([]uint8, size)
-                return &buf
-            },
-        }
-        bpm.pools[size] = pool
-    }
-    bpm.mu.Unlock()
-    return pool.Get().(*[]uint8)
+func (bpm *BufferPoolManager) GetBuffer(size uint) *[]uint8 {
+	bpm.mu.RLock()
+	pool, exists := bpm.pools[size]
+	bpm.mu.RUnlock()
+
+	if !exists {
+		bpm.mu.Lock()
+		// Double-check if another goroutine added the pool while we were waiting
+		pool, exists = bpm.pools[size]
+		if !exists {
+			pool = &sync.Pool{
+				New: func() interface{} {
+					buf := make([]uint8, size)
+					return &buf
+				},
+			}
+			bpm.pools[size] = pool
+		}
+		bpm.mu.Unlock()
+	}
+
+	return pool.Get().(*[]uint8)
 }
 
-func (bpm *BufferPoolManager) PutBuffer(buf *([]uint8)) {
-    size := uint(len(*buf))
-    bpm.mu.Lock()
-    pool, exists := bpm.pools[size]
-    if !exists {
-        bpm.mu.Unlock()
-        return
-    }
-    bpm.mu.Unlock()
-    
-    for i := range *buf {
-        (*buf)[i] = 0
-    }
-    pool.Put(buf)
+func (bpm *BufferPoolManager) PutBuffer(buf *[]uint8) {
+	size := uint(len(*buf))
+	bpm.mu.RLock()
+	pool, exists := bpm.pools[size]
+	bpm.mu.RUnlock()
+
+	if !exists {
+		return
+	}
+
+	// Clear buffer
+	for i := range *buf {
+		(*buf)[i] = 0
+	}
+
+	pool.Put(buf)
 }