Add AsyncStop config option to facilitate graceful stop in async mode on Close()

Signed-off-by: JamesJJ <jj@fcg.fyi>

Author: JamesJJ

README.md

@@ -69,6 +69,11 @@ Since the default is zero value, Write will not time out.
 Enable asynchronous I/O (connect and write) for sending events to Fluentd.
 The default is false.
 
+### AsyncStop
+
+Enables discarding bufferred events when Close() is called in Async mode. If Fluentd is continuously unavailable, Close() will block forever otherwise.
+The default is false.
+
 ### RequestAck
 
 Sets whether to request acknowledgment from Fluentd to increase the reliability

fluent/fluent.go

@@ -49,6 +49,7 @@ type Config struct {
 	MaxRetryWait     int           `json:"max_retry_wait"`
 	TagPrefix        string        `json:"tag_prefix"`
 	Async            bool          `json:"async"`
+	AsyncStop        bool          `json:"async_stop"`
 	// Deprecated: Use Async instead
 	AsyncConnect  bool `json:"async_connect"`
 	MarshalAsJSON bool `json:"marshal_as_json"`
@@ -83,8 +84,9 @@ type msgToSend struct {
 type Fluent struct {
 	Config
 
-	pending chan *msgToSend
-	wg      sync.WaitGroup
+	stopRunning chan bool
+	pending     chan *msgToSend
+	wg          sync.WaitGroup
 
 	muconn sync.Mutex
 	conn   net.Conn
@@ -305,6 +307,10 @@ func (f *Fluent) EncodeData(tag string, tm time.Time, message interface{}) (msg
 // Close closes the connection, waiting for pending logs to be sent
 func (f *Fluent) Close() (err error) {
 	if f.Config.Async {
+		if f.Config.AsyncStop {
+			f.stopRunning <- true
+			close(f.stopRunning)
+		}
 		close(f.pending)
 		f.wg.Wait()
 	}
@@ -347,13 +353,25 @@ func (f *Fluent) connect() (err error) {
 }
 
 func (f *Fluent) run() {
+	drainEvents := false
 	for {
+		select {
+		case stopRunning, ok := <-f.stopRunning:
+			if stopRunning || !ok {
+				drainEvents = true
+			}
+default:
+		}
 		select {
 		case entry, ok := <-f.pending:
 			if !ok {
 				f.wg.Done()
 				return
 			}
+			if drainEvents {
+				fmt.Fprintf(os.Stderr, "[%s] Unable to send logs to fluentd, discarding...\n", time.Now().Format(time.RFC3339))
+				continue
+			}
 			err := f.write(entry)
 			if err != nil {
 				fmt.Fprintf(os.Stderr, "[%s] Unable to send logs to fluentd, reconnecting...\n", time.Now().Format(time.RFC3339))

fluent/fluent_test.go

@@ -260,6 +260,7 @@ func TestJsonConfig(t *testing.T) {
 		MaxRetry:         3,
 		TagPrefix:        "fluent",
 		Async:            false,
+		AsyncStop:        false,
 		MarshalAsJSON:    true,
 	}
 

fluent/testdata/config.json

@@ -10,5 +10,6 @@
   "max_retry":3,
   "tag_prefix":"fluent",
   "async": false,
+  "async_stop": false,
   "marshal_as_json": true
 }