Refactor bufpool methods

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) Get(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) Put(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)
 }

pdu/bufpool_test.go

@@ -6,13 +6,10 @@ import (
 )
 
 func TestRetrieveBufferOfRequestedSize(t *testing.T) {
-	bpm := &BufferPoolManager{
-		pools: make(map[uint]*sync.Pool),
-		mu:    sync.Mutex{},
-	}
+	bpm := NewBufferPoolManager()
 
 	size := 1024
-	buffer := bpm.GetBuffer(uint(size))
+	buffer := bpm.Get(uint(size))
 
 	if buffer == nil {
 		t.Fatalf("Expected buffer, got nil")
@@ -24,13 +21,10 @@ func TestRetrieveBufferOfRequestedSize(t *testing.T) {
 }
 
 func TestRequestBufferSizeZero(t *testing.T) {
-	bpm := &BufferPoolManager{
-		pools: make(map[uint]*sync.Pool),
-		mu:    sync.Mutex{},
-	}
+	bpm := NewBufferPoolManager()
 
 	size := 0
-	buffer := bpm.GetBuffer(uint(size))
+	buffer := bpm.Get(uint(size))
 
 	if buffer == nil {
 		t.Fatalf("Expected buffer, got nil")
@@ -42,10 +36,7 @@ func TestRequestBufferSizeZero(t *testing.T) {
 }
 
 func TestConcurrentAccessToBufferPool(t *testing.T) {
-	bpm := &BufferPoolManager{
-		pools: make(map[uint]*sync.Pool),
-		mu:    sync.Mutex{},
-	}
+	bpm := NewBufferPoolManager()
 
 	size := 1024
 	var wg sync.WaitGroup
@@ -55,7 +46,7 @@ func TestConcurrentAccessToBufferPool(t *testing.T) {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			buffer := bpm.GetBuffer(uint(size))
+			buffer := bpm.Get(uint(size))
 			if buffer == nil {
 				t.Errorf("Expected buffer, got nil")
 			}
@@ -69,13 +60,10 @@ func TestConcurrentAccessToBufferPool(t *testing.T) {
 }
 
 func TestGetBufferLockUnlock(t *testing.T) {
-	bpm := &BufferPoolManager{
-		pools: make(map[uint]*sync.Pool),
-		mu:    sync.Mutex{},
-	}
+	bpm := NewBufferPoolManager()
 
 	size := 1024
-	buffer := bpm.GetBuffer(uint(size))
+	buffer := bpm.Get(uint(size))
 
 	if buffer == nil {
 		t.Fatalf("Expected buffer, got nil")
@@ -87,13 +75,10 @@ func TestGetBufferLockUnlock(t *testing.T) {
 }
 
 func TestVerifyPoolCreationForNewSizes(t *testing.T) {
-	bpm := &BufferPoolManager{
-		pools: make(map[uint]*sync.Pool),
-		mu:    sync.Mutex{},
-	}
+	bpm := NewBufferPoolManager()
 
 	size := 512
-	buffer := bpm.GetBuffer(uint(size))
+	buffer := bpm.Get(uint(size))
 
 	if buffer == nil {
 		t.Fatalf("Expected buffer, got nil")
@@ -105,13 +90,10 @@ func TestVerifyPoolCreationForNewSizes(t *testing.T) {
 }
 
 func TestBufferPoolManagerGetBuffer(t *testing.T) {
-	bpm := &BufferPoolManager{
-		pools: make(map[uint]*sync.Pool),
-		mu:    sync.Mutex{},
-	}
+	bpm := NewBufferPoolManager()
 
 	size := 1024
-	buffer := bpm.GetBuffer(uint(size))
+	buffer := bpm.Get(uint(size))
 
 	if buffer == nil {
 		t.Fatalf("Expected buffer, got nil")
@@ -123,14 +105,11 @@ func TestBufferPoolManagerGetBuffer(t *testing.T) {
 }
 
 func TestGetBufferWithMultipleSizes(t *testing.T) {
-	bpm := &BufferPoolManager{
-		pools: make(map[uint]*sync.Pool),
-		mu:    sync.Mutex{},
-	}
+	bpm := NewBufferPoolManager()
 
 	sizes := []int{512, 1024, 2048}
 	for _, size := range sizes {
-		buffer := bpm.GetBuffer(uint(size))
+		buffer := bpm.Get(uint(size))
 
 		if buffer == nil {
 			t.Fatalf("Expected buffer for size %d, got nil", size)
@@ -141,3 +120,28 @@ func TestGetBufferWithMultipleSizes(t *testing.T) {
 		}
 	}
 }
+
+func TestGetBufferIsAlwaysZero(t *testing.T) {
+	bpm := NewBufferPoolManager()
+
+	var size uint = 1024 * 64
+	for i := 0; i < 1000; i++ {
+		buffer := bpm.Get(size)
+
+		if buffer == nil {
+			t.Fatalf("Expected buffer for size %d, got nil", size)
+		}
+
+		if uint(len(*buffer)) != size {
+			t.Errorf("Expected buffer size %d, got %d", size, len(*buffer))
+		}
+
+		for _, b := range *buffer {
+			if b != 0 {
+				t.Errorf("Expected buffer to be zero, got %d", b)
+			}
+		}
+
+		bpm.Put(buffer)
+	}
+}

pdu/global.go

@@ -1,8 +1,3 @@
 package pdu
 
-import "sync"
-
-var ByteBufferPool = &BufferPoolManager{
-	pools: make(map[uint]*sync.Pool),
-	mu:    sync.Mutex{},
-}
+var ByteBufferPool = NewBufferPoolManager()

pdu/pdu.go

@@ -25,6 +25,7 @@ type (
 		header PDU_HEADER
 	}
 )
+
 // Hmm the header can be partially encoded
 // As in command_status can just be null
 // So not 0s, null, non existent, not encoded
@@ -56,9 +57,8 @@ type (
 // What are the other 0s?
 // Don't know
 
-
 func (p *PDU_HEADER) Encode() (*[]uint8, error) {
-	buf := ByteBufferPool.GetBuffer(uint(p.Size()))
+	buf := ByteBufferPool.Get(uint(p.Size()))
 	err := p.EncodeInto(buf)
 	return buf, err
 }

pdu/submit.go

@@ -61,7 +61,7 @@ type (
 )
 
 func (p *SUBMIT_SM) Encode() (*[]byte, error) {
-	buf := ByteBufferPool.GetBuffer(uint(p.Size()))
+	buf := ByteBufferPool.Get(uint(p.Size()))
 	err := p.EncodeInto(buf)
 	return buf, err
 }