Use valyala/bytebufferpool for spooledTempFile (#141)

* working

* removed old code comments

---------

Co-authored-by: Jake L <NGTmeaty@users.noreply.github.com>

Author: equals215

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/spf13/cobra v1.9.1
 	github.com/things-go/go-socks5 v0.0.6
 	github.com/ulikunitz/xz v0.5.14
+	github.com/valyala/bytebufferpool v1.0.0
 	github.com/zeebo/blake3 v0.2.4
 	go.uber.org/goleak v1.3.0
 	golang.org/x/net v0.39.0

go.sum

@@ -40,6 +40,8 @@ github.com/things-go/go-socks5 v0.0.6 h1:YjylIYZiND41szH4NzsVbx8aVDsS/Y8ps3QYPwQ
 github.com/things-go/go-socks5 v0.0.6/go.mod h1:RF6tRutwNWzISbPfiDEChH/o1aDfRv+cXDYn2a2qkK4=
 github.com/ulikunitz/xz v0.5.14 h1:uv/0Bq533iFdnMHZdRBTOlaNMdb1+ZxXIlHDZHIHcvg=
 github.com/ulikunitz/xz v0.5.14/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14=
+github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
+github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
 github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU=
 github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
 github.com/zeebo/assert v1.1.0 h1:hU1L1vLTHsnO8x8c9KAR5GmM5QscxHg5RNU5z5qbUWY=

pkg/spooledtempfile/spooled.go

@@ -9,6 +9,8 @@ import (
 	"os"
 	"sync"
 	"time"
+
+	"github.com/valyala/bytebufferpool"
 )
 
 const (
@@ -30,30 +32,8 @@ type globalMemoryCache struct {
 
 var (
 	memoryUsageCache = &globalMemoryCache{}
-	spooledPool      = sync.Pool{
-		New: func() any {
-			b := make([]byte, 0, InitialBufferSize)
-			return &b
-		},
-	}
 )
 
-// Get a zero-length slice backed by a pooled array (cap == InitialBufferSize).
-func getPooledBuf() []byte {
-	p := spooledPool.Get().(*[]byte)
-	return (*p)[:0]
-}
-
-// Return a slice to the pool without retaining large arrays.
-// IMPORTANT: do NOT pass &s.buf directly; copy a header first.
-func putPooledBuf(b []byte) {
-	if cap(b) == 0 || cap(b) > InitialBufferSize {
-		return // don't retain big arrays
-	}
-	h := b[:0]          // fresh header decoupled from callers
-	spooledPool.Put(&h) // pointer-like -> SA6002 satisfied
-}
-
 // ReaderAt is the interface for ReadAt - read at position, without moving pointer.
 type ReaderAt interface {
 	ReadAt(p []byte, off int64) (n int, err error)
@@ -70,7 +50,7 @@ type ReadSeekCloser interface {
 // spooledTempFile writes to memory (or to disk if
 // over MaxInMemorySize) and deletes the file on Close
 type spooledTempFile struct {
-	buf                 []byte        // Use []byte instead of bytes.Buffer
+	buf                 *bytebufferpool.ByteBuffer
 	mem                 *bytes.Reader // Reader for in-memory data
 	file                *os.File
 	filePrefix          string
@@ -118,7 +98,7 @@ func NewSpooledTempFile(filePrefix string, tempDir string, threshold int, fullOn
 	return &spooledTempFile{
 		filePrefix:          filePrefix,
 		tempDir:             tempDir,
-		buf:                 getPooledBuf(),
+		buf:                 bytebufferpool.Get(),
 		maxInMemorySize:     threshold,
 		fullOnDisk:          fullOnDisk,
 		maxRAMUsageFraction: maxRAMUsageFraction,
@@ -142,7 +122,7 @@ func (s *spooledTempFile) prepareRead() error {
 		return nil
 	}
 
-	s.mem = bytes.NewReader(s.buf) // Create a reader from the []byte slice
+	s.mem = bytes.NewReader(s.buf.Bytes())
 	return nil
 }
 
@@ -154,7 +134,7 @@ func (s *spooledTempFile) Len() int {
 		}
 		return int(fi.Size())
 	}
-	return len(s.buf) // Return the length of the []byte slice
+	return s.buf.Len()
 }
 
 func (s *spooledTempFile) Read(p []byte) (n int, err error) {
@@ -196,6 +176,7 @@ func (s *spooledTempFile) Write(p []byte) (n int, err error) {
 	if s.closed {
 		return 0, io.EOF
 	}
+
 	if s.reading {
 		panic("write after read")
 	}
@@ -205,27 +186,26 @@ func (s *spooledTempFile) Write(p []byte) (n int, err error) {
 	}
 
 	aboveRAMThreshold := s.isSystemMemoryUsageHigh()
-	if aboveRAMThreshold || s.fullOnDisk || (len(s.buf)+len(p) > s.maxInMemorySize) {
+	if aboveRAMThreshold || s.fullOnDisk || (s.buf.Len()+len(p) > s.maxInMemorySize) {
 		// Switch to file if we haven't already
 		s.file, err = os.CreateTemp(s.tempDir, s.filePrefix+"-")
 		if err != nil {
 			return 0, err
 		}
 
 		// Copy what we already had in the buffer
-		_, err = s.file.Write(s.buf)
+		_, err = s.buf.WriteTo(s.file)
 		if err != nil {
 			s.file.Close()
 			s.file = nil
 			return 0, err
 		}
 
 		// Release the buffer back to the pool
-		if s.buf != nil && cap(s.buf) <= InitialBufferSize && cap(s.buf) > 0 {
-			putPooledBuf(s.buf)
+		if s.buf != nil {
+			bytebufferpool.Put(s.buf)
 		}
 		s.buf = nil
-		s.mem = nil // Discard the bytes.Reader
 
 		// Write incoming bytes directly to file
 		n, err = s.file.Write(p)
@@ -237,36 +217,22 @@ func (s *spooledTempFile) Write(p []byte) (n int, err error) {
 		return n, nil
 	}
 
-	// Grow the buffer if necessary, but never exceed MaxInMemorySize
-	if len(s.buf)+len(p) > cap(s.buf) {
-		newCap := min(len(s.buf)+len(p), s.maxInMemorySize)
-
-		// Allocate a new buffer with the increased capacity
-		newBuf := make([]byte, len(s.buf), newCap)
-		copy(newBuf, s.buf)
-
-		// Release the old buffer to the pool
-		if s.buf != nil && cap(s.buf) <= InitialBufferSize && cap(s.buf) > 0 {
-			putPooledBuf(s.buf)
-		}
-		s.buf = newBuf
-		s.mem = nil // Discard the old bytes.Reader
-	}
-
 	// Append data to the buffer
-	s.buf = append(s.buf, p...)
+	s.buf.Write(p)
 	return len(p), nil
 }
 
 func (s *spooledTempFile) Close() error {
 	s.closed = true
-	s.mem = nil
+
+	if s.mem != nil {
+		s.mem.Reset([]byte{})
+		s.mem = nil
+	}
 
 	// Release the buffer back to the pool
 	if s.buf != nil {
-		if cap(s.buf) <= InitialBufferSize && cap(s.buf) > 0 {
-			putPooledBuf(s.buf)
-		}
+		bytebufferpool.Put(s.buf)
 		s.buf = nil
 	}
 

pkg/spooledtempfile/spooled_test.go

@@ -461,8 +461,8 @@ func TestBufferGrowthWithinLimits(t *testing.T) {
 
 	// Check that the buffer grew
 	spoolBuffer := spool.(*spooledTempFile)
-	if cap(spoolBuffer.buf) <= InitialBufferSize {
-		t.Fatalf("Expected buffer capacity > %d, got %d", InitialBufferSize, cap(spoolBuffer.buf))
+	if spoolBuffer.buf.Len() <= InitialBufferSize {
+		t.Fatalf("Expected buffer capacity > %d, got %d", InitialBufferSize, spoolBuffer.buf.Len())
 	}
 
 	// Check that the buffer is still in memory and has grown
@@ -492,28 +492,15 @@ func TestPoolBehavior(t *testing.T) {
 
 	// Ensure the buffer has grown beyond InitialBufferSize
 	spoolTempFile := spool.(*spooledTempFile)
-	if cap(spoolTempFile.buf) <= InitialBufferSize {
-		t.Fatalf("Expected buffer capacity > %d, got %d", InitialBufferSize, cap(spoolTempFile.buf))
+	if spoolTempFile.buf.Len() <= InitialBufferSize {
+		t.Fatalf("Expected buffer capacity > %d, got %d", InitialBufferSize, spoolTempFile.buf.Len())
 	}
 
 	// Close the spool to release the buffer
 	err = spool.Close()
 	if err != nil {
 		t.Fatalf("Close error: %v", err)
 	}
-
-	// Retrieve a buffer from the pool
-	buf := getPooledBuf()
-
-	// Verify that the retrieved buffer has the expected initial capacity
-	if cap(buf) != InitialBufferSize {
-		t.Errorf("Expected buffer in pool to have capacity %d, got %d", InitialBufferSize, cap(buf))
-	}
-
-	// Verify that the buffer is empty (reset)
-	if len(buf) != 0 {
-		t.Errorf("Expected buffer length to be 0, got %d", len(buf))
-	}
 }
 
 func TestBufferGrowthBeyondNewCap(t *testing.T) {
@@ -557,15 +544,6 @@ func TestBufferGrowthBeyondNewCap(t *testing.T) {
 	if !bytes.Equal(out, expected) {
 		t.Errorf("Data mismatch. Got %q, want %q", out, expected)
 	}
-
-	// Verify that the buffer was released to the pool (if it meets the criteria)
-	buf := getPooledBuf()
-	if cap(buf) != InitialBufferSize {
-		t.Errorf("Expected buffer in pool to have capacity %d, got %d", InitialBufferSize, cap(buf))
-	}
-	if len(buf) != 0 {
-		t.Errorf("Expected buffer length to be 0, got %d", len(buf))
-	}
 }
 
 func TestSpoolingWhenIOCopy(t *testing.T) {