diff --git a/pkg/writers/buffered_file_writer/bufferedfilewriter.go b/pkg/writers/buffered_file_writer/bufferedfilewriter.go index e5458d26b23b..7f62e81952f6 100644 --- a/pkg/writers/buffered_file_writer/bufferedfilewriter.go +++ b/pkg/writers/buffered_file_writer/bufferedfilewriter.go @@ -1,3 +1,5 @@ +// Package bufferedfilewriter provides a writer that buffers data in memory until a threshold is exceeded at +// which point it switches to writing to a temporary file. package bufferedfilewriter import ( @@ -11,9 +13,11 @@ import ( // BufferedFileWriter manages a buffer for writing data, flushing to a file when a threshold is exceeded. type BufferedFileWriter struct { - threshold uint64 - buf bytes.Buffer - file *os.File + threshold uint64 // Threshold for switching to file writing. + buf bytes.Buffer // Buffer for storing data under the threshold in memory. + filename string // Name of the temporary file. + file io.WriteCloser // File for storing data over the threshold. + size uint64 // Total size of the data written. } // Option is a function that modifies a BufferedFileWriter. @@ -41,26 +45,40 @@ func (w *BufferedFileWriter) Len() int { return w.buf.Len() } func (w *BufferedFileWriter) String() string { return w.buf.String() } // Write writes data to the buffer or a file, depending on the size. -func (w *BufferedFileWriter) Write(ctx context.Context, p []byte) (int, error) { - if uint64(w.buf.Len()+len(p)) <= w.threshold { +func (w *BufferedFileWriter) Write(ctx context.Context, data []byte) (int, error) { + size := uint64(len(data)) + defer func() { + w.size += size + ctx.Logger().V(4).Info( + "write complete", + "data_size", size, + "content_size", w.buf.Len(), + "total_size", w.size, + ) + }() + + if uint64(w.buf.Len())+size <= w.threshold { // If the total size is within the threshold, write to the buffer. ctx.Logger().V(4).Info( "writing to buffer", - "data_size", len(p), + "data_size", size, "content_size", w.buf.Len(), ) - return w.buf.Write(p) + return w.buf.Write(data) } // Switch to file writing if threshold is exceeded. // This helps in managing memory efficiently for large content. if w.file == nil { var err error - w.file, err = os.CreateTemp(os.TempDir(), cleantemp.MkFilename()) + filename := cleantemp.MkFilename() + w.file, err = os.CreateTemp(os.TempDir(), filename) if err != nil { return 0, err } + w.filename = filename + // 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 w.buf.Len() > 0 { @@ -72,9 +90,9 @@ func (w *BufferedFileWriter) Write(ctx context.Context, p []byte) (int, error) { w.buf = bytes.Buffer{} } } - ctx.Logger().V(4).Info("writing to file", "data_size", len(p)) + ctx.Logger().V(4).Info("writing to file", "data_size", size) - return w.file.Write(p) + return w.file.Write(data) } // Close flushes any remaining data in the buffer to the file and closes the file if it was created. @@ -99,7 +117,7 @@ func (w *BufferedFileWriter) Close() error { func (w *BufferedFileWriter) ReadCloser() (io.ReadCloser, error) { if w.file != nil { // Data is in a file, read from the file. - file, err := os.Open(w.file.Name()) + file, err := os.Open(w.filename) if err != nil { return nil, err } diff --git a/pkg/writers/buffered_file_writer/bufferedfilewriter_test.go b/pkg/writers/buffered_file_writer/bufferedfilewriter_test.go index 95a2190946c1..445fcb88d1f4 100644 --- a/pkg/writers/buffered_file_writer/bufferedfilewriter_test.go +++ b/pkg/writers/buffered_file_writer/bufferedfilewriter_test.go @@ -122,7 +122,7 @@ func TestBufferedFileWriterWriteExceedsThreshold(t *testing.T) { assert.NotNil(t, writer.file) assert.Len(t, writer.buf.Bytes(), 0) - fileContents, err := os.ReadFile(writer.file.Name()) + fileContents, err := os.ReadFile(writer.filename) assert.NoError(t, err) assert.Equal(t, data, fileContents) } @@ -144,9 +144,9 @@ func TestBufferedFileWriterWriteAfterFlush(t *testing.T) { defer writer.Close() // Get the file modification time after the initial write. - initialModTime, err := getFileModTime(t, writer.file.Name()) + initialModTime, err := getFileModTime(t, writer.filename) assert.NoError(t, err) - fileContents, err := os.ReadFile(writer.file.Name()) + fileContents, err := os.ReadFile(writer.filename) assert.NoError(t, err) assert.Equal(t, initialData, fileContents) @@ -155,7 +155,7 @@ func TestBufferedFileWriterWriteAfterFlush(t *testing.T) { assert.NoError(t, err) assert.Equal(t, subsequentData, writer.buf.Bytes()) // Check buffer contents - finalModTime, err := getFileModTime(t, writer.file.Name()) + finalModTime, err := getFileModTime(t, writer.filename) assert.NoError(t, err) assert.Equal(t, initialModTime, finalModTime) // File should not be modified again } @@ -229,7 +229,7 @@ func TestBufferedFileWriterClose(t *testing.T) { assert.NoError(t, err) if writer.file != nil { - fileContents, err := os.ReadFile(writer.file.Name()) + fileContents, err := os.ReadFile(writer.filename) assert.NoError(t, err) assert.Equal(t, tc.expectFileContent, string(fileContents)) return