trufflesecurity/trufflehog
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

merge

Browse Source
This commit is contained in:
Ahrav Dutta
2024-04-19 18:06:08 -07:00
parent 9dc9e15034
commit 8874bb85a4
5 changed files with 56 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/writers/buffer/buffer.go
+10 -10
View File
@@ -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
pkg/writers/buffer/buffer_test.go
-2
View File
@@ -5,8 +5,6 @@ import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
)
func TestNewBufferPool(t *testing.T) {
pkg/writers/buffer/ring/ring.go
+1 -3
View File
@@ -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
}
pkg/writers/buffer/ring/ring_test.go
+7 -12
View File
@@ -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)
pkg/writers/buffer_writer/bufferwriter_test.go
+38 -5
View File
@@ -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()
}
})
}