Skip to content

Commit

Permalink
[bug] - Correctly return the checked out buffer to the pool (#2732)
Browse files Browse the repository at this point in the history
* Make sure to return the buffer to the pool

* update comment

* defer the return

* remove anonymous function
  • Loading branch information
ahrav committed Apr 23, 2024
1 parent 46e0da1 commit ea4d9d2
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 0 deletions.
4 changes: 4 additions & 0 deletions pkg/writers/buffered_file_writer/bufferedfilewriter.go
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,10 @@ func (w *BufferedFileWriter) CloseForWriting() error {
return nil
}

// Return the buffer to the pool since the contents have been written to the file and
// the writer is transitioning to read-only mode.
defer w.bufPool.Put(w.buf)

if w.buf.Len() > 0 {
_, err := w.buf.WriteTo(w.file)
if err != nil {
Expand Down
31 changes: 31 additions & 0 deletions pkg/writers/buffered_file_writer/bufferedfilewriter_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (
"github.com/stretchr/testify/assert"

"github.com/trufflesecurity/trufflehog/v3/pkg/context"
"github.com/trufflesecurity/trufflehog/v3/pkg/writers/buffer"
)

func TestBufferedFileWriterNewThreshold(t *testing.T) {
Expand Down Expand Up @@ -506,3 +507,33 @@ func BenchmarkBufferedFileWriterWriteSmall(b *testing.B) {
rc.Close()
}
}

func TestBufferWriterCloseForWritingWithFile(t *testing.T) {
bufPool := buffer.NewBufferPool()

ctx := context.Background()
buf := bufPool.Get(ctx)
writer := &BufferedFileWriter{
threshold: 10,
bufPool: bufPool,
buf: buf,
}

// Write data exceeding the threshold to ensure a file is created.
data := []byte("this is a longer string exceeding the threshold")
_, err := writer.Write(data)
assert.NoError(t, err)

err = writer.CloseForWriting()
assert.NoError(t, err)
assert.Equal(t, readOnly, writer.state)

rdr, err := writer.ReadCloser()
assert.NoError(t, err)
defer rdr.Close()

// Get a buffer from the pool and check if it is the same buffer used in the writer.
bufFromPool := bufPool.Get(ctx)
assert.Same(t, buf, bufFromPool, "Buffer should be returned to the pool")
bufPool.Put(bufFromPool)
}

0 comments on commit ea4d9d2

Please sign in to comment.