From 72a677c91e5d10574b06b88bc15b6467ccb48bfe Mon Sep 17 00:00:00 2001
From: Ahrav Dutta <ahrav.dutta@trufflesec.com>
Date: Thu, 18 Jan 2024 07:29:40 -0800
Subject: [PATCH] use buffer pool

---
 .../bufferedfilewriter.go                     | 63 ++++++++++++++++---
 1 file changed, 53 insertions(+), 10 deletions(-)

diff --git a/pkg/writers/buffered_file_writer/bufferedfilewriter.go b/pkg/writers/buffered_file_writer/bufferedfilewriter.go
index 7f62e81952f6..02c92e8fdd50 100644
--- a/pkg/writers/buffered_file_writer/bufferedfilewriter.go
+++ b/pkg/writers/buffered_file_writer/bufferedfilewriter.go
@@ -4,13 +4,20 @@ package bufferedfilewriter
 
 import (
 	"bytes"
+	"fmt"
 	"io"
 	"os"
+	"sync"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/cleantemp"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 )
 
+// bufferPool is used to store buffers for reuse.
+var bufferPool = sync.Pool{
+	New: func() any { return new(bytes.Buffer) },
+}
+
 // BufferedFileWriter manages a buffer for writing data, flushing to a file when a threshold is exceeded.
 type BufferedFileWriter struct {
 	threshold uint64         // Threshold for switching to file writing.
@@ -57,6 +64,16 @@ func (w *BufferedFileWriter) Write(ctx context.Context, data []byte) (int, error
 		)
 	}()
 
+	if w.buf.Len() == 0 {
+		bufPtr, ok := bufferPool.Get().(*bytes.Buffer)
+		if !ok {
+			ctx.Logger().Error(fmt.Errorf("buffer pool returned unexpected type"), "using new buffer")
+			bufPtr = new(bytes.Buffer)
+		}
+		bufPtr.Reset() // Reset the buffer to clear any existing data
+		w.buf = *bufPtr
+	}
+
 	if uint64(w.buf.Len())+size <= w.threshold {
 		// If the total size is within the threshold, write to the buffer.
 		ctx.Logger().V(4).Info(
@@ -70,14 +87,13 @@ func (w *BufferedFileWriter) Write(ctx context.Context, data []byte) (int, error
 	// Switch to file writing if threshold is exceeded.
 	// This helps in managing memory efficiently for large content.
 	if w.file == nil {
-		var err error
-		filename := cleantemp.MkFilename()
-		w.file, err = os.CreateTemp(os.TempDir(), filename)
+		file, err := os.CreateTemp(os.TempDir(), cleantemp.MkFilename())
 		if err != nil {
 			return 0, err
 		}
 
-		w.filename = filename
+		w.filename = file.Name()
+		w.file = file
 
 		// Transfer existing data in buffer to the file, then clear the buffer.
 		// This ensures all the data is in one place - either entirely in the buffer or the file.
@@ -86,8 +102,9 @@ func (w *BufferedFileWriter) Write(ctx context.Context, data []byte) (int, error
 			if _, err := w.file.Write(w.buf.Bytes()); err != nil {
 				return 0, err
 			}
-			// Replace the buffer with a new one to free up memory.
-			w.buf = bytes.Buffer{}
+			// Reset the buffer to clear any existing data and return it to the pool.
+			w.buf.Reset()
+			bufferPool.Put(&w.buf)
 		}
 	}
 	ctx.Logger().V(4).Info("writing to file", "data_size", size)
@@ -110,9 +127,11 @@ func (w *BufferedFileWriter) Close() error {
 	return w.file.Close()
 }
 
-// ReadCloser returns an io.ReadCloser to read the written content. If the total content size exceeds the
-// predefined threshold, it is stored in a temporary file and a file reader is returned.
-// For content under the threshold, it is kept in memory and a bytes reader on the buffer is returned.
+// ReadCloser returns an io.ReadCloser to read the written content. It provides a reader
+// based on the current storage medium of the data (in-memory buffer or file).
+// If the total content size exceeds the predefined threshold, it is stored in a temporary file and a file
+// reader is returned. For in-memory data, it returns a custom reader that handles returning
+// the buffer to the pool.
 // The caller should call Close() on the returned io.Reader when done to ensure files are cleaned up.
 func (w *BufferedFileWriter) ReadCloser() (io.ReadCloser, error) {
 	if w.file != nil {
@@ -123,8 +142,12 @@ func (w *BufferedFileWriter) ReadCloser() (io.ReadCloser, error) {
 		}
 		return newAutoDeletingFileReader(file), nil
 	}
+
 	// Data is in memory.
-	return io.NopCloser(bytes.NewReader(w.buf.Bytes())), nil
+	return &bufferReadCloser{
+		Reader:  bytes.NewReader(w.buf.Bytes()),
+		onClose: func() { bufferPool.Put(&w.buf) },
+	}, nil
 }
 
 // autoDeletingFileReader wraps an *os.File and deletes the file on Close.
@@ -140,3 +163,23 @@ func (r *autoDeletingFileReader) Close() error {
 	defer os.Remove(r.Name()) // Delete the file after closing
 	return r.File.Close()
 }
+
+// bufferReadCloser is a custom implementation of io.ReadCloser. It wraps a bytes.Reader
+// for reading data from an in-memory buffer and includes an onClose callback.
+// The onClose callback is used to return the buffer to the pool, ensuring buffer re-usability.
+type bufferReadCloser struct {
+	*bytes.Reader
+	onClose func()
+}
+
+// Close implements the io.Closer interface. It calls the onClose callback to return the buffer
+// to the pool, enabling buffer reuse. This method should be called by the consumers of ReadCloser
+// once they have finished reading the data to ensure proper resource management.
+func (brc *bufferReadCloser) Close() error {
+	if brc.onClose == nil {
+		return nil
+	}
+
+	brc.onClose() // Return the buffer to the pool
+	return nil
+}