From 71b943b17daf9478350b1b80b15fc9dd718bca52 Mon Sep 17 00:00:00 2001
From: Kiall Mac Innes <kiall@macinnes.ie>
Date: Wed, 20 Sep 2023 10:54:54 +0000
Subject: [PATCH] Implement io.WriteCloser to allow flushing logs

---
 config.go | 14 +++++++++++---
 fluent.go | 30 ++++++++++++++++++++++++------
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/config.go b/config.go
index a53d39e..23e8885 100644
--- a/config.go
+++ b/config.go
@@ -4,9 +4,10 @@ import "github.com/fluent/fluent-logger-golang/fluent"
 
 // Config is used to configure the Writer.
 type Config struct {
-	fluentConfig *fluent.Config
-	basicConfig  *BasicConfig
-	tag          string
+	fluentConfig            *fluent.Config
+	basicConfig             *BasicConfig
+	tag                     string
+	discardWritesAfterClose bool
 }
 
 // BasicConfig is used to configure the Writer with a basic configuration.
@@ -40,3 +41,10 @@ func WithTag(tag string) func(*Config) {
 		c.tag = tag
 	}
 }
+
+// WithDiscardWritesAfterClose returns a function that can be used to configure the Writer to discard writes after Close()
+func WithDiscardWritesAfterClose() func(*Config) {
+	return func(c *Config) {
+		c.discardWritesAfterClose = true
+	}
+}
diff --git a/fluent.go b/fluent.go
index 9bec4df..f151bf6 100644
--- a/fluent.go
+++ b/fluent.go
@@ -2,17 +2,20 @@ package fluentio
 
 import (
 	"encoding/json"
-	"github.com/fluent/fluent-logger-golang/fluent"
 	"io"
+
+	"github.com/fluent/fluent-logger-golang/fluent"
 )
 
 // Writer is an io.Writer that writes to fluentd.
 type Writer struct {
-	client *fluent.Fluent
-	tag    string
+	config        *Config
+	client        *fluent.Fluent
+	tag           string
+	discardWrites bool
 }
 
-var _ io.Writer = (*Writer)(nil)
+var _ io.WriteCloser = (*Writer)(nil)
 
 // New creates a new Writer.
 // It accepts a variadic number of options that can be used to configure the Writer.
@@ -47,13 +50,20 @@ func New(opts ...func(config *Config)) (*Writer, error) {
 	}
 
 	return &Writer{
-		client: client,
-		tag:    tag,
+		config:        config,
+		client:        client,
+		tag:           tag,
+		discardWrites: false,
 	}, nil
 }
 
 // Write is the implementation of io.Writer.
 func (f *Writer) Write(p []byte) (n int, err error) {
+	// If the Writer has been closed,
+	if f.discardWrites {
+		return len(p), nil
+	}
+
 	var m map[string]interface{}
 	err = json.Unmarshal(p, &m)
 	if err != nil {
@@ -67,3 +77,11 @@ func (f *Writer) Write(p []byte) (n int, err error) {
 
 	return len(p), nil
 }
+
+// Close is the implementation of io.Closer.
+func (f *Writer) Close() error {
+	if f.config.discardWritesAfterClose {
+		f.discardWrites = true
+	}
+	return f.client.Close()
+}