From bccc11ac2101aa4265612e8a72447402d4909d49 Mon Sep 17 00:00:00 2001 From: Ahrav Dutta Date: Tue, 30 Apr 2024 19:52:10 -0700 Subject: [PATCH] use ring buffer --- pkg/buffers/buffer/buffer.go | 73 +++- pkg/buffers/buffer/buffer_test.go | 82 ++++- pkg/buffers/buffer/metrics.go | 4 +- pkg/buffers/buffer/ring/ring.go | 215 ++++++++++++ pkg/buffers/buffer/ring/ring_test.go | 331 ++++++++++++++++++ pkg/buffers/pool/pool.go | 37 +- pkg/buffers/pool/pool_test.go | 12 +- pkg/handlers/handlers_test.go | 10 +- pkg/writers/buffer_writer/bufferwriter.go | 11 +- .../bufferedfilewriter.go | 10 +- 10 files changed, 714 insertions(+), 71 deletions(-) create mode 100644 pkg/buffers/buffer/ring/ring.go create mode 100644 pkg/buffers/buffer/ring/ring_test.go diff --git a/pkg/buffers/buffer/buffer.go b/pkg/buffers/buffer/buffer.go index dba11fe35..c0cd52ba7 100644 --- a/pkg/buffers/buffer/buffer.go +++ b/pkg/buffers/buffer/buffer.go @@ -4,27 +4,62 @@ package buffer import ( "bytes" + "io" "time" + + "github.com/trufflesecurity/trufflehog/v3/pkg/buffers/buffer/ring" ) -// Buffer is a wrapper around bytes.Buffer that includes a timestamp for tracking Buffer checkout duration. -type Buffer struct { - *bytes.Buffer +type Buffer interface { + io.ReadWriter + io.WriterTo + Len() int + Cap() int + Reset() + Grow(int) + Bytes() []byte + String() string +} + +// CheckoutBuffer is a wrapper around bytes.Buffer that includes a timestamp for tracking CheckoutBuffer checkout duration. +type CheckoutBuffer struct { + Buffer checkedOutAt time.Time } const defaultBufferSize = 1 << 12 // 4KB -// NewBuffer creates a new instance of Buffer. -func NewBuffer() *Buffer { return &Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, defaultBufferSize))} } -func (b *Buffer) Grow(size int) { +// WithBuffer is a functional option to specify a custom buffer. +func WithBuffer(b Buffer) func(*CheckoutBuffer) { + return func(cb *CheckoutBuffer) { + cb.Buffer = b + } +} + +// Option is a functional option to configure a CheckoutBuffer. +type Option func(*CheckoutBuffer) + +// NewBuffer creates a new instance of CheckoutBuffer. +func NewBuffer(opts ...Option) *CheckoutBuffer { + buf := new(CheckoutBuffer) + for _, opt := range opts { + opt(buf) + } + if buf.Buffer == nil { + buf.Buffer = ring.NewBuffer(defaultBufferSize) + } + + return buf +} + +func (b *CheckoutBuffer) Grow(size int) { b.Buffer.Grow(size) b.recordGrowth(size) } -func (b *Buffer) ResetMetric() { b.checkedOutAt = time.Now() } +func (b *CheckoutBuffer) ResetMetric() { b.checkedOutAt = time.Now() } -func (b *Buffer) RecordMetric() { +func (b *CheckoutBuffer) RecordMetric() { dur := time.Since(b.checkedOutAt) checkoutDuration.Observe(float64(dur.Microseconds())) checkoutDurationTotal.Add(float64(dur.Microseconds())) @@ -32,16 +67,16 @@ func (b *Buffer) RecordMetric() { totalBufferLength.Add(float64(b.Len())) } -func (b *Buffer) recordGrowth(size int) { +func (b *CheckoutBuffer) recordGrowth(size int) { growCount.Inc() growAmount.Add(float64(size)) } -// Write date to the buffer. -func (b *Buffer) Write(data []byte) (int, error) { +// Write date to the Buffer. +func (b *CheckoutBuffer) Write(data []byte) (int, error) { if b.Buffer == nil { // This case should ideally never occur if buffers are properly managed. - b.Buffer = bytes.NewBuffer(make([]byte, 0, defaultBufferSize)) + b.Buffer = ring.NewBuffer(defaultBufferSize) b.ResetMetric() } @@ -49,11 +84,11 @@ func (b *Buffer) Write(data []byte) (int, error) { bufferLength := b.Buffer.Len() totalSizeNeeded := bufferLength + size - // If the total size is within the threshold, write to the buffer. + // If the total size is within the threshold, write to the Buffer. availableSpace := b.Buffer.Cap() - bufferLength growSize := totalSizeNeeded - bufferLength if growSize > availableSpace { - // We are manually growing the buffer so we can track the growth via metrics. + // We are manually growing the Buffer so we can track the growth via metrics. // Knowing the exact data size, we directly resize to fit it, rather than exponential growth // 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 @@ -65,8 +100,8 @@ func (b *Buffer) Write(data []byte) (int, error) { } // readCloser 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. +// 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 readCloser struct { *bytes.Reader onClose func() @@ -77,14 +112,14 @@ func ReadCloser(data []byte, onClose func()) *readCloser { return &readCloser{Reader: bytes.NewReader(data), onClose: onClose} } -// 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 +// 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 *readCloser) Close() error { if brc.onClose == nil { return nil } - brc.onClose() // Return the buffer to the pool + brc.onClose() // Return the Buffer to the pool return nil } diff --git a/pkg/buffers/buffer/buffer_test.go b/pkg/buffers/buffer/buffer_test.go index bbf4987e2..f22babde9 100644 --- a/pkg/buffers/buffer/buffer_test.go +++ b/pkg/buffers/buffer/buffer_test.go @@ -5,9 +5,11 @@ import ( "testing" "github.com/stretchr/testify/assert" + + "github.com/trufflesecurity/trufflehog/v3/pkg/buffers/buffer/ring" ) -func TestBufferWrite(t *testing.T) { +func TestBytesBufferWrite(t *testing.T) { t.Parallel() tests := []struct { name string @@ -17,7 +19,7 @@ func TestBufferWrite(t *testing.T) { expectedCap int }{ { - name: "Write to empty buffer", + name: "Write to empty Buffer", initialCapacity: defaultBufferSize, writeDataSequence: [][]byte{ []byte("hello"), @@ -66,7 +68,81 @@ func TestBufferWrite(t *testing.T) { tc := tc t.Run(tc.name, func(t *testing.T) { t.Parallel() - buf := &Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, tc.initialCapacity))} + buf := &CheckoutBuffer{Buffer: bytes.NewBuffer(make([]byte, 0, tc.initialCapacity))} + totalWritten := 0 + for _, data := range tc.writeDataSequence { + n, err := buf.Write(data) + assert.NoError(t, err) + + totalWritten += n + } + assert.Equal(t, tc.expectedSize, totalWritten) + assert.Equal(t, tc.expectedSize, buf.Len()) + assert.GreaterOrEqual(t, buf.Cap(), tc.expectedCap) + }) + } +} + +func TestRingBufferWrite(t *testing.T) { + t.Parallel() + tests := []struct { + name string + initialCapacity int + writeDataSequence [][]byte // Sequence of writes to simulate multiple writes + expectedSize int + expectedCap int + }{ + { + name: "Write to empty Buffer", + initialCapacity: defaultBufferSize, + writeDataSequence: [][]byte{ + []byte("hello"), + }, + expectedSize: 5, + expectedCap: defaultBufferSize, // No growth for small data + }, + { + name: "Write causing growth", + initialCapacity: 10, // Small initial capacity to force growth + writeDataSequence: [][]byte{ + []byte("this is a longer string exceeding initial capacity"), + }, + expectedSize: 50, + expectedCap: 50, + }, + { + name: "Write nil data", + initialCapacity: defaultBufferSize, + writeDataSequence: [][]byte{nil}, + expectedCap: defaultBufferSize, + }, + { + name: "Repeated writes, cumulative growth", + initialCapacity: 20, // Set an initial capacity to test growth over multiple writes + writeDataSequence: [][]byte{ + []byte("first write, "), + []byte("second write, "), + []byte("third write exceeding the initial capacity."), + }, + expectedSize: 70, + expectedCap: 70, // Expect capacity to grow to accommodate all writes + }, + { + name: "Write large single data to test significant growth", + initialCapacity: 50, // Set an initial capacity smaller than the data to be written + writeDataSequence: [][]byte{ + bytes.Repeat([]byte("a"), 1024), // 1KB data to significantly exceed initial capacity + }, + expectedSize: 1024, + expectedCap: 1024, + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + buf := &CheckoutBuffer{Buffer: ring.NewBuffer(tc.initialCapacity)} totalWritten := 0 for _, data := range tc.writeDataSequence { n, err := buf.Write(data) diff --git a/pkg/buffers/buffer/metrics.go b/pkg/buffers/buffer/metrics.go index 4a22580eb..6358f7ffa 100644 --- a/pkg/buffers/buffer/metrics.go +++ b/pkg/buffers/buffer/metrics.go @@ -26,14 +26,14 @@ var ( Namespace: common.MetricsNamespace, Subsystem: common.MetricsSubsystem, Name: "checkout_duration_total_us", - Help: "Total duration in microseconds of Buffer checkouts.", + Help: "Total duration in microseconds of CheckoutBuffer checkouts.", }) checkoutDuration = promauto.NewHistogram(prometheus.HistogramOpts{ Namespace: common.MetricsNamespace, Subsystem: common.MetricsSubsystem, Name: "checkout_duration_us", - Help: "Duration in microseconds of Buffer checkouts.", + Help: "Duration in microseconds of CheckoutBuffer checkouts.", Buckets: prometheus.ExponentialBuckets(10, 10, 7), }) diff --git a/pkg/buffers/buffer/ring/ring.go b/pkg/buffers/buffer/ring/ring.go new file mode 100644 index 000000000..755f7854c --- /dev/null +++ b/pkg/buffers/buffer/ring/ring.go @@ -0,0 +1,215 @@ +package ring + +import ( + "io" +) + +// Ring is a ring buffer implementation that implements the io.Writer and io.Reader interfaces. +type Ring struct { + buf []byte + size int + rp int // Read pointer + wp int // Write pointer + isFull bool + // availableCap tracks the available capacity in the buffer. + // It avoids recalculating the available capacity on every write operation, + // which is needed in case the buffer needs to be resized. + availableCap int +} + +const defaultBufferSize = 1 << 12 // 4KB + +// NewBuffer creates a new ring buffer with the given size. +func NewBuffer(size int) *Ring { + return &Ring{ + buf: make([]byte, size), + size: size, + availableCap: size, + } +} + +// Write writes len(p) bytes from p to the underlying data buffer. +func (r *Ring) Write(p []byte) (int, error) { + if len(p) == 0 { + return 0, nil + } + + if len(p) > r.availableCap { + oldLen := r.Len() + newSize := r.size * 2 + if newSize < oldLen+len(p) { + newSize = oldLen + len(p) + // Add extra capacity to ensure the buffer is not resized to fit only the new data. + newSize += newSize / 2 + } + r.resize(newSize, oldLen) // Ensure there's enough space for new data + } + n := r.write(p) + + return n, nil +} + +// WriteTo writes the buffer's data to w. +func (r *Ring) WriteTo(w io.Writer) (int64, error) { + if r.isEmpty() { + return 0, nil // Nothing to write + } + + var totalWritten int64 + + // If the data wraps in the buffer, write in two parts: from read pointer to end, and from start to write pointer. + if r.wp <= r.rp { + // Data does not wrap, can directly write from r to w + n, err := w.Write(r.buf[r.rp:r.wp]) + totalWritten += int64(n) + if err != nil { + return totalWritten, err + } + r.rp = 0 + } + + // Data does not wrap, or we are continuing from a previous write after a wrap. + if r.wp > 0 { + // Data wraps, write from r to end of buffer. + n, err := w.Write(r.buf[r.rp:]) + totalWritten += int64(n) + if err != nil { + r.rp = (r.rp + n) % r.size + return totalWritten, err + } + } + + // Update the ring buffer's read pointer and full flag after successful write. + // Since all data has been read, we can reset the ring buffer. + r.Reset() + + return totalWritten, nil +} + +func (r *Ring) write(p []byte) int { + endPos := (r.wp + len(p)) % r.size + n := len(p) // The number of bytes to write is determined by the length of p. + + // Execute the copy operation based on the relative positions of wp and rp. + if r.wp >= r.rp { + // Write wraps the buffer when the buffer is being re-used. + if endPos < r.wp { + copy(r.buf[r.wp:], p[:r.size-r.wp]) + copy(r.buf[:endPos], p[r.size-r.wp:]) + } else { + copy(r.buf[r.wp:], p) + } + } else { + copy(r.buf[r.wp:endPos], p) + } + + r.wp = endPos + // The buffer is full if after the write, the next write position (wp) would overwrite the read position (rp). + // Calculate the next write position after this write. + nextWp := (r.wp + 1) % r.size + // Initially check if the next position after wp is rp, indicating a full buffer on the next write. + r.isFull = nextWp == r.rp + + // Exactly full if the next write position is the read position. + if r.wp == r.rp && n > 0 { + r.isFull = true + } + + r.availableCap -= n + if r.availableCap < 0 { + r.availableCap = 0 + } + + return n +} + +func (r *Ring) resize(newSize, oldLen int) { + newBuf := make([]byte, newSize) + r.read(newBuf) + r.buf = newBuf + r.size = newSize + r.rp = 0 + r.wp = oldLen + r.isFull = false + + r.availableCap = newSize - oldLen +} + +// Read reads up to len(p) bytes into p from the underlying data buffer. +func (r *Ring) Read(p []byte) (int, error) { + if r.isEmpty() { + return 0, nil + } + + var n int + if r.wp > r.rp { + // Non-wrapped data: directly copy from rp to wp. + n = copy(p, r.buf[r.rp:r.wp]) + } else { + // Wrapped data: copy in two steps. + // First, from rp to the end of the buffer. + n = copy(p, r.buf[r.rp:]) + // If there's more space in p and data wrapped, copy the remaining data from the beginning. + if n < len(p) { + n += copy(p[n:], r.buf[:r.wp]) + } + } + + // Update the read pointer and re-calculate the available capacity. + r.rp = (r.rp + n) % r.size + r.isFull = false // The buffer cannot be full after a read operation. + + r.availableCap = r.size - r.Len() + + return n, nil +} + +func (r *Ring) read(p []byte) int { + n := copy(p, r.buf[r.rp:]) + if r.wp < r.rp { + n += copy(p[n:], r.buf[:r.wp]) + } + return n +} + +// Bytes returns a copy of the buffer's data. +func (r *Ring) Bytes() []byte { + if r.isEmpty() { + return nil + } + + return r.buf[r.rp:r.wp] +} + +// String returns the buffer's data as a string. +func (r *Ring) String() string { return string(r.Bytes()) } + +// isEmpty returns true if the buffer is empty. +func (r *Ring) isEmpty() bool { return !r.isFull && r.rp == r.wp } + +// Len return the length of available read bytes. +func (r *Ring) Len() int { + if r.isFull { + return r.size + } + if r.wp >= r.rp { + return r.wp - r.rp + } + return r.size - r.rp + r.wp +} + +// Cap returns the size of the underlying buffer. +func (r *Ring) Cap() int { + return r.size +} + +// Reset resets the buffer. +func (r *Ring) Reset() { + r.rp = 0 + r.wp = 0 + r.isFull = false + r.availableCap = r.size +} + +// Grow is a no-op for ring buffer. +func (*Ring) Grow(_ int) {} diff --git a/pkg/buffers/buffer/ring/ring_test.go b/pkg/buffers/buffer/ring/ring_test.go new file mode 100644 index 000000000..e62e84c0e --- /dev/null +++ b/pkg/buffers/buffer/ring/ring_test.go @@ -0,0 +1,331 @@ +package ring + +import ( + "math/rand" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestRingWrite(t *testing.T) { + type writeTest struct { + name string + initBufSize int // Initial buffer size + initData []byte // Initial data to preload the buffer + writeData []byte // Data to write in the test + wantWritten int // Expected number of bytes written + wantBufSize int // Expected buffer size after write + wantBufData []byte // Expected buffer data after write + } + + testCases := []writeTest{ + { + name: "WriteWithEmptyBuffer", + initBufSize: 10, + initData: nil, + writeData: []byte("hello"), + wantWritten: 5, + wantBufSize: 10, + wantBufData: []byte("hello"), + }, + { + name: "WriteWithExactFillNoInitialData", + initBufSize: 5, + initData: nil, + writeData: []byte("hello"), + wantWritten: 5, + wantBufSize: 5, + wantBufData: []byte("hello"), + }, + { + name: "WriteWithPartialFill", + initBufSize: 10, + initData: []byte("hi"), + writeData: []byte("there"), + wantWritten: 5, + wantBufSize: 10, + wantBufData: []byte("hithere"), + }, + { + name: "WriteWithExactFill", + initBufSize: 5, + initData: []byte("hi"), + writeData: []byte("123"), + wantWritten: 3, + wantBufSize: 5, + wantBufData: []byte("hi123"), + }, + { + name: "WriteWithOverflowRequiresResize", + initBufSize: 5, + initData: []byte("hi"), + writeData: []byte("12345"), + wantWritten: 5, + wantBufSize: 10, // Expected double the size of the initial buffer + wantBufData: []byte("hi12345\x00\x00\x00"), // Including the extra null bytes + }, + { + name: "WriteWithZeroLengthData", + initBufSize: 10, + initData: []byte("data"), + writeData: nil, + wantWritten: 0, + wantBufSize: 10, + wantBufData: []byte("data"), + }, + } + + for _, tc := range testCases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + r := NewBuffer(tc.initBufSize) + if len(tc.initData) > 0 { + _, _ = r.Write(tc.initData) + } + + gotWritten, err := r.Write(tc.writeData) + assert.NoError(t, err) + + assert.Equal(t, tc.wantWritten, gotWritten) + assert.Equal(t, tc.wantBufSize, len(r.buf)) + if tc.wantWritten >= tc.initBufSize { + assert.Equal(t, tc.wantBufData, r.buf) + } else { + assert.Equal(t, tc.wantBufData, r.buf[:tc.wantWritten+len(tc.initData)]) + } + }) + } +} + +func BenchmarkBufferWrite(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 := NewBuffer(defaultBufferSize) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, err := r.Write(data) + assert.NoError(b, err) + r.Reset() + } + }) + } +} + +func generateData(size int) []byte { + rand.Seed(42) + data := make([]byte, size) + for i := range data { + data[i] = byte(rand.Intn(256)) + } + return data +} + +func TestRingBufferBytes(t *testing.T) { + tests := []struct { + name string + initBufSize int + initData []byte + wantBytes []byte + }{ + { + name: "EmptyBuffer", + initBufSize: 10, + }, + { + name: "PartialFill", + initBufSize: 10, + initData: []byte("hello"), + wantBytes: []byte("hello"), + }, + { + name: "FullFill", + initBufSize: 5, + initData: []byte("hello"), + wantBytes: []byte("hello"), + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + r := NewBuffer(tc.initBufSize) + if tc.initData != nil { + _, _ = r.Write(tc.initData) + } + + assert.Equal(t, tc.wantBytes, r.Bytes()) + }) + } +} + +func TestRingBufferLen(t *testing.T) { + tests := []struct { + name string + initBufSize int + initData []byte + wantLen int + }{ + { + name: "EmptyBuffer", + initBufSize: 10, + }, + { + name: "PartialFill", + initBufSize: 10, + initData: []byte("hello"), + wantLen: 5, + }, + { + name: "FullFill", + initBufSize: 5, + initData: []byte("hello"), + wantLen: 5, + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + r := NewBuffer(tc.initBufSize) + if tc.initData != nil { + _, _ = r.Write(tc.initData) + } + + assert.Equal(t, tc.wantLen, r.Len()) + }) + } +} + +func TestRingBufferCap(t *testing.T) { + tests := []struct { + name string + initBufSize int + wantCap int + }{ + { + name: "InitialCapacity", + initBufSize: 10, + wantCap: 10, + }, + { + name: "ZeroCapacity", + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + r := NewBuffer(tc.initBufSize) + + assert.Equal(t, tc.wantCap, r.Cap()) + }) + } +} + +func TestRingBufferReset(t *testing.T) { + tests := []struct { + name string + initBufSize int + initData []byte + }{ + { + name: "AfterWrite", + initBufSize: 10, + initData: []byte("reset"), + }, + { + name: "EmptyBuffer", + initBufSize: 10, + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + r := NewBuffer(tc.initBufSize) + if tc.initData != nil { + _, _ = r.Write(tc.initData) + } + + r.Reset() + assert.Equal(t, 0, r.Len()) + assert.Equal(t, tc.initBufSize, r.Cap()) + assert.True(t, r.isEmpty()) + }) + } +} + +func TestRingBufferRead(t *testing.T) { + tests := []struct { + name string + initBufSize int + writeData []byte + readLen int + wantReadData []byte + }{ + { + name: "ReadPartialData", + initBufSize: 10, + writeData: []byte("hello"), + readLen: 5, + wantReadData: []byte("hello"), + }, + { + name: "ReadFullData", + initBufSize: 5, + writeData: []byte("hello"), + readLen: 5, + wantReadData: []byte("hello"), + }, + { + name: "ReadEmptyBuffer", + initBufSize: 10, + writeData: nil, + readLen: 5, + wantReadData: []byte{}, + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + r := NewBuffer(tc.initBufSize) + if tc.writeData != nil { + _, _ = r.Write(tc.writeData) + } + + readData := make([]byte, tc.readLen) + n, err := r.Read(readData) + assert.NoError(t, err) + + assert.Equal(t, tc.wantReadData, readData[:n]) + }) + } +} diff --git a/pkg/buffers/pool/pool.go b/pkg/buffers/pool/pool.go index 9696f5bf2..fdb153941 100644 --- a/pkg/buffers/pool/pool.go +++ b/pkg/buffers/pool/pool.go @@ -1,7 +1,6 @@ package pool import ( - "bytes" "sync" "github.com/trufflesecurity/trufflehog/v3/pkg/buffers/buffer" @@ -20,18 +19,14 @@ func (poolMetrics) recordBufferRetrival() { bufferCount.Inc() } -func (poolMetrics) recordBufferReturn(buf *buffer.Buffer) { +func (poolMetrics) recordBufferReturn(buf *buffer.CheckoutBuffer) { activeBufferCount.Dec() buf.RecordMetric() } -// Opts is a function that configures a BufferPool. -type Opts func(pool *Pool) - // Pool of buffers. type Pool struct { *sync.Pool - bufferSize uint32 metrics poolMetrics } @@ -39,26 +34,22 @@ type Pool struct { // const defaultBufferSize = 1 << 12 // 4KB const defaultBufferSize = 1 << 16 // 64KB // NewBufferPool creates a new instance of BufferPool. -func NewBufferPool(opts ...Opts) *Pool { - pool := &Pool{bufferSize: defaultBufferSize} +func NewBufferPool(opts ...buffer.Option) *Pool { + pool := new(Pool) - for _, opt := range opts { - opt(pool) - } + buf := buffer.NewBuffer(opts...) pool.Pool = &sync.Pool{ - New: func() any { - return &buffer.Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, pool.bufferSize))} - }, + New: func() any { return buf }, } return pool } -// Get returns a Buffer from the pool. -func (p *Pool) Get() *buffer.Buffer { - buf, ok := p.Pool.Get().(*buffer.Buffer) +// Get returns a CheckoutBuffer from the pool. +func (p *Pool) Get() *buffer.CheckoutBuffer { + buf, ok := p.Pool.Get().(*buffer.CheckoutBuffer) if !ok { - buf = &buffer.Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, p.bufferSize))} + buf = buffer.NewBuffer() } p.metrics.recordBufferRetrival() buf.ResetMetric() @@ -66,18 +57,18 @@ func (p *Pool) Get() *buffer.Buffer { return buf } -// Put returns a Buffer to the pool. -func (p *Pool) Put(buf *buffer.Buffer) { +// Put returns a CheckoutBuffer to the pool. +func (p *Pool) Put(buf *buffer.CheckoutBuffer) { p.metrics.recordBufferReturn(buf) - // If the Buffer is more than twice the default size, replace it with a new Buffer. + // If the CheckoutBuffer is more than twice the default size, replace it with a new CheckoutBuffer. // This prevents us from returning very large buffers to the pool. const maxAllowedCapacity = 2 * defaultBufferSize if buf.Cap() > maxAllowedCapacity { p.metrics.recordShrink(buf.Cap() - defaultBufferSize) - buf = &buffer.Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, p.bufferSize))} + buf = buffer.NewBuffer() } else { - // Reset the Buffer to clear any existing data. + // Reset the CheckoutBuffer to clear any existing data. buf.Reset() } diff --git a/pkg/buffers/pool/pool_test.go b/pkg/buffers/pool/pool_test.go index 1968356e1..1c2c0d1f5 100644 --- a/pkg/buffers/pool/pool_test.go +++ b/pkg/buffers/pool/pool_test.go @@ -38,13 +38,13 @@ func TestBufferPoolGetPut(t *testing.T) { t.Parallel() tests := []struct { name string - preparePool func(p *Pool) *buffer.Buffer // Prepare the pool and return an initial buffer to put if needed - expectedCapBefore int // Expected capacity before putting it back - expectedCapAfter int // Expected capacity after retrieving it again + preparePool func(p *Pool) *buffer.CheckoutBuffer // Prepare the pool and return an initial buffer to put if needed + expectedCapBefore int // Expected capacity before putting it back + expectedCapAfter int // Expected capacity after retrieving it again }{ { name: "Get new buffer and put back without modification", - preparePool: func(_ *Pool) *buffer.Buffer { + preparePool: func(_ *Pool) *buffer.CheckoutBuffer { return nil // No initial buffer to put }, expectedCapBefore: int(defaultBufferSize), @@ -52,8 +52,8 @@ func TestBufferPoolGetPut(t *testing.T) { }, { name: "Put oversized buffer, expect shrink", - preparePool: func(p *Pool) *buffer.Buffer { - buf := &buffer.Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, 3*defaultBufferSize))} + preparePool: func(p *Pool) *buffer.CheckoutBuffer { + buf := &buffer.CheckoutBuffer{Buffer: bytes.NewBuffer(make([]byte, 0, 3*defaultBufferSize))} return buf }, expectedCapBefore: int(defaultBufferSize), diff --git a/pkg/handlers/handlers_test.go b/pkg/handlers/handlers_test.go index 52082f7fa..fd65d05ae 100644 --- a/pkg/handlers/handlers_test.go +++ b/pkg/handlers/handlers_test.go @@ -1,7 +1,6 @@ package handlers import ( - "io" "net/http" "os" "strings" @@ -40,12 +39,10 @@ func TestHandleFile(t *testing.T) { } func BenchmarkHandleFile(b *testing.B) { - file, err := os.Open("testdata/test.tgz") - assert.Nil(b, err) - - b.ResetTimer() for i := 0; i < b.N; i++ { sourceChan := make(chan *sources.Chunk, 1) + file, err := os.Open("testdata/test.tgz") + assert.Nil(b, err) b.StartTimer() go func() { @@ -57,9 +54,6 @@ func BenchmarkHandleFile(b *testing.B) { for range sourceChan { } b.StopTimer() - - _, err = file.Seek(0, io.SeekStart) - assert.NoError(b, err) } } diff --git a/pkg/writers/buffer_writer/bufferwriter.go b/pkg/writers/buffer_writer/bufferwriter.go index a522acca0..513a0b462 100644 --- a/pkg/writers/buffer_writer/bufferwriter.go +++ b/pkg/writers/buffer_writer/bufferwriter.go @@ -20,7 +20,7 @@ func (metrics) recordDataProcessed(size int64, dur time.Duration) { func init() { bufferPool = pool.NewBufferPool() } -// bufferPool is the shared Buffer pool used by all BufferedFileWriters. +// bufferPool is the shared CheckoutBuffer pool used by all BufferedFileWriters. // This allows for efficient reuse of buffers across multiple writers. var bufferPool *pool.Pool @@ -36,10 +36,10 @@ const ( // BufferWriter implements contentWriter, using a shared buffer pool for memory management. type BufferWriter struct { - buf *buffer.Buffer // The current buffer in use. - bufPool *pool.Pool // The buffer pool used to manage the buffer. - size int // The total size of the content written to the buffer. - state state // The current state of the buffer. + buf *buffer.CheckoutBuffer // The current buffer in use. + bufPool *pool.Pool // The buffer pool used to manage the buffer. + size int // The total size of the content written to the buffer. + state state // The current state of the buffer. metrics metrics } @@ -62,6 +62,7 @@ func NewFromReader(ctx context.Context, r io.Reader) (*BufferWriter, error) { return buf, nil } +const defaultBufferSize = 1 << 12 // 4KB // Write delegates the writing operation to the underlying bytes.Buffer. func (b *BufferWriter) Write(data []byte) (int, error) { if b.state != writeOnly { diff --git a/pkg/writers/buffered_file_writer/bufferedfilewriter.go b/pkg/writers/buffered_file_writer/bufferedfilewriter.go index f6232155c..0cee91a2f 100644 --- a/pkg/writers/buffered_file_writer/bufferedfilewriter.go +++ b/pkg/writers/buffered_file_writer/bufferedfilewriter.go @@ -30,7 +30,7 @@ func (bufferedFileWriterMetrics) recordDiskWrite(size int64) { func init() { bufferPool = pool.NewBufferPool() } -// bufferPool is the shared Buffer pool used by all BufferedFileWriters. +// bufferPool is the shared CheckoutBuffer pool used by all BufferedFileWriters. // This allows for efficient reuse of buffers across multiple writers. var bufferPool *pool.Pool @@ -50,10 +50,10 @@ type BufferedFileWriter struct { threshold uint64 // Threshold for switching to file writing. size uint64 // Total size of the data written. - bufPool *pool.Pool // Pool for storing buffers for reuse. - buf *buffer.Buffer // Buffer for storing data under the threshold in memory. - filename string // Name of the temporary file. - file *os.File // File for storing data over the threshold. + bufPool *pool.Pool // Pool for storing buffers for reuse. + buf *buffer.CheckoutBuffer // CheckoutBuffer for storing data under the threshold in memory. + filename string // Name of the temporary file. + file *os.File // File for storing data over the threshold. state state // Current state of the writer. (writeOnly or readOnly)