feat: Apply upload criteria to memory buffering; Upload the memory buffer to S3 on exit.

internal/irzstd/disk.go

@@ -260,14 +260,6 @@ func (w *diskWriter) Close() error {
 	return nil
 }
 
-// Getter for useDiskBuffer.
-//
-// Returns:
-//   - useDiskBuffer: On/off for disk buffering
-func (w *diskWriter) GetUseDiskBuffer() bool {
-	return true
-}
-
 // Getter for Zstd Output.
 //
 // Returns:
@@ -284,6 +276,21 @@ func (w *diskWriter) GetZstdOutputSize() (int, error) {
 	return w.getZstdFileSize()
 }
 
+// Checks if writer is empty. True if no events are buffered.
+//
+// Returns:
+//   - empty: Boolean value that is true if buffer is empty
+//   - err: Error calling stat
+func (w *diskWriter) CheckEmpty() (bool, error) {
+	zstdFileInfo, err := w.zstdFile.Stat()
+	if err != nil {
+		return false, err
+	}
+
+	empty := (zstdFileInfo.Size() == 0) && (w.irTotalBytes == 0)
+	return empty, nil
+}
+
 // Compresses contents of the IR file and outputs it to the Zstd file. The IR file is then
 // truncated.
 //

internal/irzstd/memory.go

@@ -11,9 +11,8 @@ import (
 	"github.com/y-scope/clp-ffi-go/ir"
 )
 
-// Converts log events into Zstd compressed IR. Log events provided to writer are immediately
-// converted to Zstd compressed IR and stored in [memoryWriter.ZstdBuffer].  After the Zstd buffer
-// receives logs, they are immediately sent to s3.
+// Converts log events into Zstd compressed IR. Log events are immediately converted to Zstd
+// compressed IR and stored in [memoryWriter.zstdBuffer].
 type memoryWriter struct {
 	zstdBuffer *bytes.Buffer
 	irWriter   *ir.Writer
@@ -96,14 +95,6 @@ func (w *memoryWriter) Reset() error {
 	return nil
 }
 
-// Getter for useDiskBuffer.
-//
-// Returns:
-//   - useDiskBuffer: On/off for disk buffering
-func (w *memoryWriter) GetUseDiskBuffer() bool {
-	return false
-}
-
 // Getter for Zstd Output.
 //
 // Returns:
@@ -113,7 +104,8 @@ func (w *memoryWriter) GetZstdOutput() io.Reader {
 }
 
 // Get size of Zstd output. [zstd] does not provide the amount of bytes written with each write.
-// Instead, calling Len() on buffer.
+// Instead, calling Len() on buffer. Size may slightly lag the real size since some data in the
+// current block will be in the [zstd] encoder's internal buffer.
 //
 // Returns:
 //   - size: Bytes written
@@ -139,3 +131,16 @@ func (w *memoryWriter) Close() error {
 	}
 	return nil
 }
+
+// Checks if writer is empty. True if no events are buffered. Try to avoid calling this as will
+// flush Zstd Writer potentially creating unnecessary frames.
+//
+// Returns:
+//   - empty: Boolean value that is true if buffer is empty
+//   - err: nil error to comply with interface
+func (w *memoryWriter) CheckEmpty() (bool, error) {
+	w.zstdWriter.Flush()
+
+	empty := w.zstdBuffer.Len() == 0
+	return empty, nil
+}

internal/irzstd/writer.go

@@ -43,12 +43,6 @@ type Writer interface {
 	//   - err
 	Reset() error
 
-	// Getter for useDiskBuffer.
-	//
-	// Returns:
-	//   - useDiskBuffer: On/off for disk buffering
-	GetUseDiskBuffer() bool
-
 	// Getter for Zstd Output.
 	//
 	// Returns:
@@ -61,6 +55,13 @@ type Writer interface {
 	//	 - size: Bytes written
 	//   - err
 	GetZstdOutputSize() (int, error)
+
+	// Checks if writer is empty. True if no events are buffered.
+	//
+	// Returns:
+	//   - empty: Boolean value that is true if buffer is empty
+	//   - err
+	CheckEmpty() (bool, error)
 }
 
 // Writes log events to a IR Writer.

plugins/out_clp_s3/internal/exit/exit.go

@@ -0,0 +1,59 @@
+// Package exit provides functions for gracefully shutting down the plugin. Exit functions are only
+// called when Fluent Bit receives a kill signal, not during an abrupt crash. The plugin is given
+// limited time to clean up resources before Fluent Bit terminates it.
+
+package exit
+
+import (
+	"github.com/y-scope/fluent-bit-clp/internal/outctx"
+)
+
+// Fs gracefully exits the plugin by closing files.
+//
+// Parameters:
+//   - ctx: Plugin context
+//
+// Returns:
+//   - err: Error closing file
+func Fs(ctx *outctx.S3Context) error {
+	for _, eventManager := range ctx.EventManagers {
+		err := eventManager.Writer.Close()
+		if err != nil {
+			return err
+		}
+		eventManager.Writer = nil
+	}
+
+	return nil
+}
+
+// S3 gracefully exits the plugin by flushing buffered data to S3. Makes a best-effort attempt,
+// however Fluent Bit may kill the plugin before the upload completes.
+//
+// Parameters:
+//   - ctx: Plugin context
+//
+// Returns:
+//   - err: Error closing file
+func S3(ctx *outctx.S3Context) error {
+	for _, eventManager := range ctx.EventManagers {
+		empty, err := eventManager.Writer.CheckEmpty()
+		if err != nil {
+			return err
+		}
+		if empty {
+			continue
+		}
+		err = eventManager.ToS3(ctx.Config, ctx.Uploader)
+		if err != nil {
+			return err
+		}
+		err = eventManager.Writer.Close()
+		if err != nil {
+			return err
+		}
+		eventManager.Writer = nil
+	}
+
+	return nil
+}

plugins/out_clp_s3/internal/flush/flush.go

@@ -111,9 +111,7 @@ func decodeMsgpack(dec *codec.Decoder, config outctx.S3Config) ([]ffi.LogEvent,
 	}
 }
 
-// Checks if criteria are met to upload to s3. If useDiskBuffer is false, then the chunk is always
-// uploaded so always returns true. If useDiskBuffer is true, check if Zstd buffer size is greater
-// than upload size.
+// Checks whether Zstd buffer size is greater than or equal to upload size.
 //
 // Parameters:
 //   - eventManager: Manager for Fluent Bit events with the same tag
@@ -123,10 +121,6 @@ func decodeMsgpack(dec *codec.Decoder, config outctx.S3Config) ([]ffi.LogEvent,
 //   - readyToUpload: Boolean if upload criteria met or not
 //   - err: Error getting Zstd buffer size
 func checkUploadCriteriaMet(eventManager *outctx.EventManager, uploadSizeMb int) (bool, error) {
-	if !eventManager.Writer.GetUseDiskBuffer() {
-		return true, nil
-	}
-
 	bufferSize, err := eventManager.Writer.GetZstdOutputSize()
 	if err != nil {
 		return false, fmt.Errorf("error could not get size of buffer: %w", err)

plugins/out_clp_s3/internal/recovery/recovery.go

@@ -13,28 +13,6 @@ import (
 	"github.com/y-scope/fluent-bit-clp/internal/outctx"
 )
 
-// If useDiskBuffer is set, close all files prior to exit. Graceful exit will only be called
-// if Fluent Bit receives a kill signal and not during an abrupt crash. Plugin is only
-// given a limited time to clean up resources, so output is not sent to s3. Instead
-// they are sent during startup.
-//
-// Parameters:
-//   - ctx: Plugin context
-//
-// Returns:
-//   - err: Error closing file
-func GracefulExit(ctx *outctx.S3Context) error {
-	for _, eventManager := range ctx.EventManagers {
-		err := eventManager.Writer.Close()
-		if err != nil {
-			return err
-		}
-		eventManager.Writer = nil
-	}
-
-	return nil
-}
-
 // Sends existing disk buffers to S3.
 //
 // Parameters:

plugins/out_clp_s3/out_clp_s3.go

@@ -16,6 +16,7 @@ import (
 	"github.com/fluent/fluent-bit-go/output"
 
 	"github.com/y-scope/fluent-bit-clp/internal/outctx"
+	"github.com/y-scope/fluent-bit-clp/plugins/out_clp_s3/internal/exit"
 	"github.com/y-scope/fluent-bit-clp/plugins/out_clp_s3/internal/flush"
 	"github.com/y-scope/fluent-bit-clp/plugins/out_clp_s3/internal/recovery"
 )
@@ -132,7 +133,12 @@ func FLBPluginExitCtx(ctx unsafe.Pointer) int {
 
 	log.Printf("[%s] Exit called for id: %s", s3PluginName, outCtx.Config.Id)
 
-	err := recovery.GracefulExit(outCtx)
+	var err error
+	if outCtx.Config.UseDiskBuffer {
+		err = exit.Fs(outCtx)
+	} else {
+		err = exit.S3(outCtx)
+	}
 	if err != nil {
 		log.Printf("Failed to exit gracefully")
 	}