diff --git a/pkg/writers/buffered_file_writer/bufferedfilewriter.go b/pkg/writers/buffered_file_writer/bufferedfilewriter.go index 63c6733bad1f..8fa93fe16ccb 100644 --- a/pkg/writers/buffered_file_writer/bufferedfilewriter.go +++ b/pkg/writers/buffered_file_writer/bufferedfilewriter.go @@ -205,7 +205,13 @@ func (w *BufferedFileWriter) Write(data []byte) (int, error) { } return 0, err } - w.buf.Reset() + } + } + + // Write any remaining data in the buffer to the file before writing new data. + if w.buf.Len() > 0 { + if _, err := w.buf.WriteTo(w.file); err != nil { + return 0, fmt.Errorf("error flushing buffer to file: %w", err) } } @@ -273,8 +279,7 @@ func (w *BufferedFileWriter) CloseForWriting() error { defer w.bufPool.Put(w.buf) if w.buf.Len() > 0 { - _, err := w.buf.WriteTo(w.file) - if err != nil { + if _, err := w.buf.WriteTo(w.file); err != nil { return err } } diff --git a/pkg/writers/buffered_file_writer/bufferedfilewriter_test.go b/pkg/writers/buffered_file_writer/bufferedfilewriter_test.go index 310b8086c855..f250f9420e96 100644 --- a/pkg/writers/buffered_file_writer/bufferedfilewriter_test.go +++ b/pkg/writers/buffered_file_writer/bufferedfilewriter_test.go @@ -650,11 +650,19 @@ func TestBufferedFileWriter_ReadFrom(t *testing.T) { expectedSize: 1 << 20, }, { - name: "Input greater than threshold", + name: "Input slightly greater than threshold", input: string(make([]byte, defaultThreshold+1)), expectedOutput: string(make([]byte, defaultThreshold+1)), expectedSize: defaultThreshold + 1, }, + // Test to ensure that anytime the buffer exceeds the threshold, the data is written to a file + // and the buffer is cleared. + { + name: "Input much greater than threshold", + input: string(make([]byte, (2*defaultThreshold)+largeBufferSize+1)), + expectedOutput: string(make([]byte, (2*defaultThreshold)+largeBufferSize+1)), + expectedSize: (2 * defaultThreshold) + largeBufferSize + 1, + }, } for _, tc := range tests { @@ -668,6 +676,10 @@ func TestBufferedFileWriter_ReadFrom(t *testing.T) { size, err := writer.ReadFrom(reader) assert.NoError(t, err) + if writer.buf != nil && writer.file != nil { + assert.Len(t, writer.buf.Bytes(), 0) + } + err = writer.CloseForWriting() assert.NoError(t, err)