diff --git a/pkg/writers/buffer/buffer.go b/pkg/writers/buffer/buffer.go index a9f4bb42c..eceab2cf0 100644 --- a/pkg/writers/buffer/buffer.go +++ b/pkg/writers/buffer/buffer.go @@ -14,20 +14,20 @@ type Buffer struct { // NewBuffer creates a new instance of Buffer. func NewBuffer() *Buffer { return &Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, defaultBufferSize))} } -func (r *Buffer) Grow(size int) { - r.Buffer.Grow(size) - r.recordGrowth(size) +func (b *Buffer) Grow(size int) { + b.Buffer.Grow(size) + b.recordGrowth(size) } -func (r *Buffer) resetMetric() { r.checkedOutAt = time.Now() } +func (b *Buffer) resetMetric() { b.checkedOutAt = time.Now() } -func (r *Buffer) recordMetric() { - dur := time.Since(r.checkedOutAt) +func (b *Buffer) recordMetric() { + dur := time.Since(b.checkedOutAt) checkoutDuration.Observe(float64(dur.Microseconds())) checkoutDurationTotal.Add(float64(dur.Microseconds())) } -func (r *Buffer) recordGrowth(size int) { +func (b *Buffer) recordGrowth(size int) { growCount.Inc() growAmount.Add(float64(size)) } @@ -41,7 +41,7 @@ func (b *Buffer) Write(data []byte) (int, error) { } size := len(data) - bufferLength := r.Buffer.Len() + bufferLength := b.Buffer.Len() totalSizeNeeded := bufferLength + size // If the total size is within the threshold, write to the buffer. @@ -53,10 +53,10 @@ func (b *Buffer) Write(data []byte) (int, error) { // which may require multiple allocations and copies if the size required is much larger // than double the capacity. Our approach aligns with default behavior when growth sizes // happen to match current capacity, retaining asymptotic efficiency benefits. - r.Buffer.Grow(growSize) + b.Buffer.Grow(growSize) } - return r.Buffer.Write(data) + return b.Buffer.Write(data) } // readCloser is a custom implementation of io.ReadCloser. It wraps a bytes.Reader diff --git a/pkg/writers/buffer/buffer_test.go b/pkg/writers/buffer/buffer_test.go index a062450f8..ea0cd5583 100644 --- a/pkg/writers/buffer/buffer_test.go +++ b/pkg/writers/buffer/buffer_test.go @@ -5,8 +5,6 @@ import ( "testing" "github.com/stretchr/testify/assert" - - "github.com/trufflesecurity/trufflehog/v3/pkg/context" ) func TestNewBufferPool(t *testing.T) { diff --git a/pkg/writers/buffer/ring/ring.go b/pkg/writers/buffer/ring/ring.go index f1719e3f5..815a0a48d 100644 --- a/pkg/writers/buffer/ring/ring.go +++ b/pkg/writers/buffer/ring/ring.go @@ -2,8 +2,6 @@ package ring import ( "io" - - "github.com/trufflesecurity/trufflehog/v3/pkg/context" ) // Ring is a ring buffer implementation that implements the io.Writer and io.Reader interfaces. @@ -31,7 +29,7 @@ func NewRingBuffer(size int) *Ring { } // Write writes len(p) bytes from p to the underlying data buffer. -func (r *Ring) Write(_ context.Context, p []byte) (int, error) { +func (r *Ring) Write(p []byte) (int, error) { if len(p) == 0 { return 0, nil } diff --git a/pkg/writers/buffer/ring/ring_test.go b/pkg/writers/buffer/ring/ring_test.go index 16d34b7cb..d9677edcb 100644 --- a/pkg/writers/buffer/ring/ring_test.go +++ b/pkg/writers/buffer/ring/ring_test.go @@ -5,8 +5,6 @@ import ( "testing" "github.com/stretchr/testify/assert" - - "github.com/trufflesecurity/trufflehog/v3/pkg/context" ) func TestRingWrite(t *testing.T) { @@ -83,12 +81,11 @@ func TestRingWrite(t *testing.T) { t.Parallel() r := NewRingBuffer(tc.initBufSize) - ctx := context.Background() if len(tc.initData) > 0 { - _, _ = r.Write(ctx, tc.initData) + _, _ = r.Write(tc.initData) } - gotWritten, err := r.Write(ctx, tc.writeData) + gotWritten, err := r.Write(tc.writeData) assert.NoError(t, err) assert.Equal(t, tc.wantWritten, gotWritten) @@ -123,14 +120,12 @@ func BenchmarkBufferWrite(b *testing.B) { for _, bc := range benchmarks { bc := bc b.Run(bc.name, func(b *testing.B) { - ctx := context.Background() - data := generateData(bc.dataSize) // Generate pseudo-random data for this benchmark case r := NewRingBuffer(defaultBufferSize) b.ResetTimer() for i := 0; i < b.N; i++ { - _, err := r.Write(ctx, data) + _, err := r.Write(data) assert.NoError(b, err) r.Reset() } @@ -178,7 +173,7 @@ func TestRingBufferBytes(t *testing.T) { t.Parallel() r := NewRingBuffer(tc.initBufSize) if tc.initData != nil { - _, _ = r.Write(context.Background(), tc.initData) + _, _ = r.Write(tc.initData) } assert.Equal(t, tc.wantBytes, r.Bytes()) @@ -217,7 +212,7 @@ func TestRingBufferLen(t *testing.T) { t.Parallel() r := NewRingBuffer(tc.initBufSize) if tc.initData != nil { - _, _ = r.Write(context.Background(), tc.initData) + _, _ = r.Write(tc.initData) } assert.Equal(t, tc.wantLen, r.Len()) @@ -275,7 +270,7 @@ func TestRingBufferReset(t *testing.T) { t.Parallel() r := NewRingBuffer(tc.initBufSize) if tc.initData != nil { - _, _ = r.Write(context.Background(), tc.initData) + _, _ = r.Write(tc.initData) } r.Reset() @@ -323,7 +318,7 @@ func TestRingBufferRead(t *testing.T) { t.Parallel() r := NewRingBuffer(tc.initBufSize) if tc.writeData != nil { - _, _ = r.Write(context.Background(), tc.writeData) + _, _ = r.Write(tc.writeData) } readData := make([]byte, tc.readLen) diff --git a/pkg/writers/buffer_writer/bufferwriter_test.go b/pkg/writers/buffer_writer/bufferwriter_test.go index 64f94740a..e59bbfb81 100644 --- a/pkg/writers/buffer_writer/bufferwriter_test.go +++ b/pkg/writers/buffer_writer/bufferwriter_test.go @@ -11,6 +11,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/trufflesecurity/trufflehog/v3/pkg/context" + "github.com/trufflesecurity/trufflehog/v3/pkg/writers/buffer" "github.com/trufflesecurity/trufflehog/v3/pkg/writers/buffer/ring" ) @@ -169,6 +170,40 @@ func generateData(size int) []byte { return data } +func BenchmarkRingBufferWrite(b *testing.B) { + type benchCase struct { + name string + dataSize int // Size of the data to write in bytes + } + + benchmarks := []benchCase{ + {"1KB", 1 << 10}, // 1KB + {"4KB", 4 << 10}, // 4KB + {"16KB", 16 << 10}, // 16KB + {"64KB", 64 << 10}, // 64KB + {"256KB", 256 << 10}, // 256KB + {"1MB", 1 << 20}, // 1MB + {"4MB", 4 << 20}, // 4MB + {"16MB", 16 << 20}, // 16MB + {"64MB", 64 << 20}, // 64MB + } + + for _, bc := range benchmarks { + bc := bc + b.Run(bc.name, func(b *testing.B) { + data := generateData(bc.dataSize) // Generate pseudo-random data for this benchmark case + r := ring.NewRingBuffer(bc.dataSize) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, err := r.Write(data) + assert.NoError(b, err) + r.Reset() + } + }) + } +} + func BenchmarkBufferWrite(b *testing.B) { type benchCase struct { name string @@ -190,16 +225,14 @@ func BenchmarkBufferWrite(b *testing.B) { for _, bc := range benchmarks { bc := bc b.Run(bc.name, func(b *testing.B) { - ctx := context.Background() - data := generateData(bc.dataSize) // Generate pseudo-random data for this benchmark case - r := ring.NewRingBuffer(bc.dataSize) + buf := buffer.NewBuffer() b.ResetTimer() for i := 0; i < b.N; i++ { - _, err := r.Write(ctx, data) + _, err := buf.Write(data) assert.NoError(b, err) - r.Reset() + buf.Reset() } }) }