Mirrors/trufflehog
2
0
Fork 0
mirror of https://github.com/trufflesecurity/trufflehog.git synced 2024-11-10 07:04:24 +00:00
Code Issues Projects Releases Packages Wiki Activity

[refactor] - Update Write method signature in contentWriter interface (#2721)

Browse source 
* Update write method in contentWriter interface

* fix lint
This commit is contained in:
ahrav 2024-04-23 08:47:53 -07:00 committed by GitHub
parent 642fce5edf
commit 4a5fbf8417
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 43 additions and 74 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
pkg/gitparse/gitparse.go
View file

@ -37,7 +37,7 @@ const (
// based on performance needs or resource constraints, providing a unified way to interact with different content types.
type contentWriter interface { // Write appends data to the content storage.
// Write appends data to the content storage.
Write(ctx context.Context, data []byte) (int, error)
Write(data []byte) (int, error)
// ReadCloser provides a reader for accessing stored content.
ReadCloser() (io.ReadCloser, error)
// CloseForWriting closes the content storage for writing.
@ -92,8 +92,8 @@ func (d *Diff) Len() int { return d.contentWriter.Len() }
func (d *Diff) ReadCloser() (io.ReadCloser, error) { return d.contentWriter.ReadCloser() }
// write delegates to the contentWriter.
func (d *Diff) write(ctx context.Context, p []byte) error {
_, err := d.contentWriter.Write(ctx, p)
func (d *Diff) write(p []byte) error {
_, err := d.contentWriter.Write(p)
return err
}
@ -483,7 +483,7 @@ func (c *Parser) FromReader(ctx context.Context, stdOut io.Reader, diffChan chan
latestState = HunkContentLine
}
// TODO: Why do we care about this? It creates empty lines in the diff. If there are no plusLines, it's just newlines.
if err := currentDiff.write(ctx, []byte("\n")); err != nil {
if err := currentDiff.write([]byte("\n")); err != nil {
ctx.Logger().Error(err, "failed to write to diff")
}
case isHunkPlusLine(latestState, line):
@ -491,7 +491,7 @@ func (c *Parser) FromReader(ctx context.Context, stdOut io.Reader, diffChan chan
latestState = HunkContentLine
}
if err := currentDiff.write(ctx, line[1:]); err != nil {
if err := currentDiff.write(line[1:]); err != nil {
ctx.Logger().Error(err, "failed to write to diff")
}
// NoOp. We only care about additions.

4
pkg/gitparse/gitparse_test.go
View file

@ -694,7 +694,7 @@ func TestCommitParsing(t *testing.T) {
func newBufferedFileWriterWithContent(content []byte) *bufferedfilewriter.BufferedFileWriter {
ctx := context.Background()
b := bufferedfilewriter.New(ctx)
_, err := b.Write(ctx, content) // Using Write method to add content
_, err := b.Write(content) // Using Write method to add content
if err != nil {
panic(err)
}
@ -704,7 +704,7 @@ func newBufferedFileWriterWithContent(content []byte) *bufferedfilewriter.Buffer
func newBufferWithContent(content []byte) *bufferwriter.BufferWriter {
ctx := context.Background()
b := bufferwriter.New(ctx)
_, _ = b.Write(ctx, content) // Using Write method to add content
_, _ = b.Write(content) // Using Write method to add content
return b
}

17
pkg/writers/buffer/buffer.go
View file

@ -118,10 +118,9 @@ func (b *Buffer) recordGrowth(size int) {
}
// Write date to the buffer.
func (b *Buffer) Write(ctx context.Context, data []byte) (int, error) {
func (b *Buffer) Write(data []byte) (int, error) {
if b.Buffer == nil {
// This case should ideally never occur if buffers are properly managed.
ctx.Logger().Error(fmt.Errorf("buffer is nil, initializing a new buffer"), "action", "initializing_new_buffer")
b.Buffer = bytes.NewBuffer(make([]byte, 0, defaultBufferSize))
b.resetMetric()
}
@ -129,23 +128,11 @@ func (b *Buffer) Write(ctx context.Context, data []byte) (int, error) {
size := len(data)
bufferLength := b.Buffer.Len()
totalSizeNeeded := bufferLength + size
// If the total size is within the threshold, write to the buffer.
ctx.Logger().V(4).Info(
"writing to buffer",
"data_size", size,
"content_size", bufferLength,
)
// If the total size is within the threshold, write to the buffer.
availableSpace := b.Buffer.Cap() - bufferLength
growSize := totalSizeNeeded - bufferLength
if growSize > availableSpace {
ctx.Logger().V(4).Info(
"buffer size exceeded, growing buffer",
"current_size", bufferLength,
"new_size", totalSizeNeeded,
"available_space", availableSpace,
"grow_size", growSize,
)
// 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

2
pkg/writers/buffer/buffer_test.go
View file

@ -144,7 +144,7 @@ func TestBufferWrite(t *testing.T) {
buf := &Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, tc.initialCapacity))}
totalWritten := 0
for _, data := range tc.writeDataSequence {
n, err := buf.Write(context.Background(), data)
n, err := buf.Write(data)
assert.NoError(t, err)
totalWritten += n

14
pkg/writers/buffer_writer/bufferwriter.go
View file

@ -52,10 +52,8 @@ func New(ctx context.Context) *BufferWriter {
return &BufferWriter{buf: buf, state: writeOnly, bufPool: bufferPool}
}
// Write delegates the writing operation to the underlying bytes.Buffer, ignoring the context.
// The context is included to satisfy the contentWriter interface, allowing for future extensions
// where context handling might be necessary (e.g., for timeouts or cancellation).
func (b *BufferWriter) Write(ctx context.Context, data []byte) (int, error) {
// Write delegates the writing operation to the underlying bytes.Buffer.
func (b *BufferWriter) Write(data []byte) (int, error) {
if b.state != writeOnly {
return 0, fmt.Errorf("buffer must be in write-only mode to write data; current state: %d", b.state)
}
@ -67,14 +65,8 @@ func (b *BufferWriter) Write(ctx context.Context, data []byte) (int, error) {
bufferLength := uint64(b.buf.Len())
b.metrics.recordDataProcessed(bufferLength, time.Since(start))
ctx.Logger().V(4).Info(
"write complete",
"data_size", size,
"buffer_len", bufferLength,
"buffer_size", b.buf.Cap(),
)
}(start)
return b.buf.Write(ctx, data)
return b.buf.Write(data)
}
// ReadCloser provides a read-closer for the buffer's content.

4
pkg/writers/buffer_writer/bufferwriter_test.go
View file

@ -44,7 +44,7 @@ func TestBufferWriterWrite(t *testing.T) {
writer := New(context.Background())
writer.state = tc.initialState
_, err := writer.Write(context.Background(), tc.input)
_, err := writer.Write(tc.input)
if tc.expectedError {
assert.Error(t, err)
} else {
@ -121,7 +121,7 @@ func TestBufferWriterString(t *testing.T) {
{
name: "String with data",
prepareBuffer: func(bw *BufferWriter) {
_, _ = bw.Write(context.Background(), []byte("test data"))
_, _ = bw.Write([]byte("test data"))
},
expectedStr: "test data",
expectedError: false,

18
pkg/writers/buffered_file_writer/bufferedfilewriter.go
View file

@ -27,11 +27,10 @@ func (bufferedFileWriterMetrics) recordDataProcessed(size uint64, dur time.Durat
totalWriteDuration.Add(float64(dur.Microseconds()))
}
func (bufferedFileWriterMetrics) recordDiskWrite(ctx context.Context, f *os.File) {
func (bufferedFileWriterMetrics) recordDiskWrite(f *os.File) {
diskWriteCount.Inc()
size, err := f.Stat()
if err != nil {
ctx.Logger().Error(err, "failed to get file size for metric")
return
}
fileSizeHistogram.Observe(float64(size.Size()))
@ -122,7 +121,7 @@ func (w *BufferedFileWriter) String() (string, error) {
}
// Write writes data to the buffer or a file, depending on the size.
func (w *BufferedFileWriter) Write(ctx context.Context, data []byte) (int, error) {
func (w *BufferedFileWriter) Write(data []byte) (int, error) {
if w.state != writeOnly {
return 0, fmt.Errorf("BufferedFileWriter must be in write-only mode to write")
}
@ -132,19 +131,12 @@ func (w *BufferedFileWriter) Write(ctx context.Context, data []byte) (int, error
start := time.Now()
defer func(start time.Time) {
w.metrics.recordDataProcessed(size, time.Since(start))
w.size += size
ctx.Logger().V(4).Info(
"write complete",
"data_size", size,
"content_size", bufferLength,
"total_size", w.size,
)
}(start)
totalSizeNeeded := uint64(bufferLength) + size
if totalSizeNeeded <= w.threshold {
return w.buf.Write(ctx, data)
return w.buf.Write(data)
}
// Switch to file writing if threshold is exceeded.
@ -157,19 +149,17 @@ func (w *BufferedFileWriter) Write(ctx context.Context, data []byte) (int, error
w.filename = file.Name()
w.file = file
w.metrics.recordDiskWrite(ctx, file)
w.metrics.recordDiskWrite(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.
if bufferLength > 0 {
ctx.Logger().V(4).Info("writing buffer to file", "content_size", bufferLength)
if _, err := w.buf.WriteTo(w.file); err != nil {
return 0, err
}
w.bufPool.Put(w.buf)
}
}
ctx.Logger().V(4).Info("writing to file", "data_size", size)
return w.file.Write(data)
}

48
pkg/writers/buffered_file_writer/bufferedfilewriter_test.go
View file

@ -79,12 +79,12 @@ func TestBufferedFileWriterString(t *testing.T) {
ctx := context.Background()
writer := New(ctx, WithThreshold(tc.threshold))
// First write, should go to file if it exceeds the threshold.
_, err := writer.Write(ctx, tc.input)
_, err := writer.Write(tc.input)
assert.NoError(t, err)
// Second write, should go to buffer
if tc.additionalInput != nil {
_, err = writer.Write(ctx, tc.additionalInput)
_, err = writer.Write(tc.additionalInput)
assert.NoError(t, err)
}
@ -111,7 +111,7 @@ func BenchmarkBufferedFileWriterString_BufferOnly_Small(b *testing.B) {
ctx := context.Background()
writer := New(ctx)
_, err := writer.Write(ctx, data)
_, err := writer.Write(data)
assert.NoError(b, err)
benchmarkBufferedFileWriterString(b, writer)
@ -129,7 +129,7 @@ func BenchmarkBufferedFileWriterString_BufferOnly_Medium(b *testing.B) {
ctx := context.Background()
writer := New(ctx)
_, err := writer.Write(ctx, data)
_, err := writer.Write(data)
assert.NoError(b, err)
benchmarkBufferedFileWriterString(b, writer)
@ -148,7 +148,7 @@ func BenchmarkBufferedFileWriterString_OnlyFile_Small(b *testing.B) {
ctx := context.Background()
writer := New(ctx)
_, err := writer.Write(ctx, data)
_, err := writer.Write(data)
assert.NoError(b, err)
benchmarkBufferedFileWriterString(b, writer)
@ -167,7 +167,7 @@ func BenchmarkBufferedFileWriterString_OnlyFile_Medium(b *testing.B) {
ctx := context.Background()
writer := New(ctx)
_, err := writer.Write(ctx, data)
_, err := writer.Write(data)
assert.NoError(b, err)
benchmarkBufferedFileWriterString(b, writer)
@ -186,11 +186,11 @@ func BenchmarkBufferedFileWriterString_BufferWithFile_Small(b *testing.B) {
ctx := context.Background()
writer := New(ctx)
_, err := writer.Write(ctx, data)
_, err := writer.Write(data)
assert.NoError(b, err)
// Write again so we also fill up the buffer.
_, err = writer.Write(ctx, data)
_, err = writer.Write(data)
assert.NoError(b, err)
benchmarkBufferedFileWriterString(b, writer)
@ -209,11 +209,11 @@ func BenchmarkBufferedFileWriterString_BufferWithFile_Medium(b *testing.B) {
ctx := context.Background()
writer := New(ctx)
_, err := writer.Write(ctx, data)
_, err := writer.Write(data)
assert.NoError(b, err)
// Write again so we also fill up the buffer.
_, err = writer.Write(ctx, data)
_, err = writer.Write(data)
assert.NoError(b, err)
benchmarkBufferedFileWriterString(b, writer)
@ -253,7 +253,7 @@ func TestBufferedFileWriterLen(t *testing.T) {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
writer := New(context.Background())
_, err := writer.Write(context.Background(), tc.input)
_, err := writer.Write(tc.input)
assert.NoError(t, err)
length := writer.Len()
@ -271,7 +271,7 @@ func TestBufferedFileWriterWriteWithinThreshold(t *testing.T) {
data := []byte("hello world")
writer := New(ctx, WithThreshold(64))
_, err := writer.Write(ctx, data)
_, err := writer.Write(data)
assert.NoError(t, err)
assert.Equal(t, data, writer.buf.Bytes())
@ -286,7 +286,7 @@ func TestBufferedFileWriterWriteExceedsThreshold(t *testing.T) {
data := []byte("hello world")
writer := New(ctx, WithThreshold(5))
_, err := writer.Write(ctx, data)
_, err := writer.Write(data)
assert.NoError(t, err)
defer func() {
@ -312,7 +312,7 @@ func TestBufferedFileWriterWriteAfterFlush(t *testing.T) {
// Initialize writer with a threshold that initialData will exceed.
writer := New(ctx, WithThreshold(uint64(len(initialData)-1)))
_, err := writer.Write(ctx, initialData)
_, err := writer.Write(initialData)
assert.NoError(t, err)
defer func() {
@ -328,7 +328,7 @@ func TestBufferedFileWriterWriteAfterFlush(t *testing.T) {
assert.Equal(t, initialData, fileContents)
// Perform a subsequent write with data under the threshold.
_, err = writer.Write(ctx, subsequentData)
_, err = writer.Write(subsequentData)
assert.NoError(t, err)
assert.Equal(t, subsequentData, writer.buf.Bytes()) // Check buffer contents
@ -362,7 +362,7 @@ func TestBufferedFileWriterClose(t *testing.T) {
name: "No File Created, Only Buffer Data",
prepareWriter: func(w *BufferedFileWriter) {
// Write data under the threshold
_, _ = w.Write(ctx, []byte("small data"))
_, _ = w.Write([]byte("small data"))
},
expectFileContent: "",
},
@ -370,7 +370,7 @@ func TestBufferedFileWriterClose(t *testing.T) {
name: "File Created, No Data in Buffer",
prepareWriter: func(w *BufferedFileWriter) {
// Write data over the threshold to create a file
_, _ = w.Write(ctx, []byte("large data is more than the threshold"))
_, _ = w.Write([]byte("large data is more than the threshold"))
},
expectFileContent: "large data is more than the threshold",
},
@ -378,8 +378,8 @@ func TestBufferedFileWriterClose(t *testing.T) {
name: "File Created, Data in Buffer",
prepareWriter: func(w *BufferedFileWriter) {
// Write data over the threshold to create a file, then write more data
_, _ = w.Write(ctx, []byte("large data is more than the threshold"))
_, _ = w.Write(ctx, []byte(" more data"))
_, _ = w.Write([]byte("large data is more than the threshold"))
_, _ = w.Write([]byte(" more data"))
},
expectFileContent: "large data is more than the threshold more data",
},
@ -387,7 +387,7 @@ func TestBufferedFileWriterClose(t *testing.T) {
name: "File Created, Buffer Cleared",
prepareWriter: func(w *BufferedFileWriter) {
// Write data over the threshold to create a file, then clear the buffer.
_, _ = w.Write(ctx, []byte("large data is more than the threshold"))
_, _ = w.Write([]byte("large data is more than the threshold"))
w.buf.Reset()
},
expectFileContent: "large data is more than the threshold",
@ -426,7 +426,7 @@ func TestBufferedFileWriterStateTransitionOnClose(t *testing.T) {
assert.Equal(t, writeOnly, writer.state)
// Perform some write operation.
_, err := writer.Write(context.Background(), []byte("test data"))
_, err := writer.Write([]byte("test data"))
assert.NoError(t, err)
// Close the writer.
@ -443,7 +443,7 @@ func TestBufferedFileWriterWriteInReadOnlyState(t *testing.T) {
_ = writer.CloseForWriting() // Transition to read-only mode
// Attempt to write in read-only mode.
_, err := writer.Write(context.Background(), []byte("should fail"))
_, err := writer.Write([]byte("should fail"))
assert.Error(t, err)
}
@ -462,7 +462,7 @@ func BenchmarkBufferedFileWriterWriteLarge(b *testing.B) {
b.StartTimer()
{
_, err := writer.Write(ctx, data)
_, err := writer.Write(data)
assert.NoError(b, err)
}
b.StopTimer()
@ -492,7 +492,7 @@ func BenchmarkBufferedFileWriterWriteSmall(b *testing.B) {
b.StartTimer()
{
_, err := writer.Write(ctx, data)
_, err := writer.Write(data)
assert.NoError(b, err)
}
b.StopTimer()