Properly stop logger during (re)connect failure

PR fluent#77 introduced a new parameter named ForceStopAsyncSend. It can be
used to tell the logger to not try to send all the log messages in its
buffer before closing. Without this parameter, the logger hangs out
whenever it has logs to write and the target Fluentd server is down.

But this new parameter is not enough: the connection is currently lazily
initialized when the logger receive its first message. This blocks the
select reading messages from the queue (until the connection is ready).
Moreover, the connection dialing uses an exponential back-off retry.

Because of that, the logger won't look for messages on `stopRunning`
channel (the channel introduced by PR fluent#77), either because it's
blocked by the Sleep used for the retry or because the connection
dialing is waiting until dialing timeout (eg. TCP timeout).

To fix these edge cases, the time.Sleep() used for back-off retry has
been transformed into a time.After(). Moreover, the dialer.Dial() call
used to connect to Fluentd has been replaced with dialer.DialContext()
and a cancellable context is now used to stop the call to that method.
The associated cancel function is stored in Fluent and got called by
Close() when ForceStopAsyncSend option is enabled.

Finally, the Fluent.run() method has been adapted to wait for both new
messages on f.pending and stop signal sent to f.stopRunning channel.
Previously, both channels were selected independently, in a procedural
fashion.

This fix is motivated by the issue described in: moby/moby#40063.

Signed-off-by: Albin Kerouanton <[email protected]>

Author: Albin Kerouanton

fluent/fluent.go

@@ -1,10 +1,12 @@
 package fluent
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"math"
+	"math/rand"
 	"net"
 	"os"
 	"reflect"
@@ -15,7 +17,6 @@ import (
 	"bytes"
 	"encoding/base64"
 	"encoding/binary"
-	"math/rand"
 
 	"github.com/tinylib/msgp/msgp"
 )
@@ -37,6 +38,13 @@ const (
 	defaultSubSecondPrecision = false
 )
 
+// randomGenerator is used by getUniqueId to generate ack hashes.
+// Unfortunately, using a seeded random number generator isn't enough to ensure
+// ack hashes are deterministically generated during tests, thus we need to
+// change randomGenerator value to ensure the hashes are stable during tests
+// such that we can expect how the logger behaves with RequestAck option enabled.
+var randomGenerator = rand.Uint64
+
 type Config struct {
 	FluentPort         int           `json:"fluent_port"`
 	FluentHost         string        `json:"fluent_host"`
@@ -85,15 +93,22 @@ type msgToSend struct {
 type Fluent struct {
 	Config
 
-	dialer      dialer
-	stopRunning chan bool
-	pending     chan *msgToSend
-	wg          sync.WaitGroup
+	dialer dialer
+	// stopRunning is used in async mode to signal to run() it should abort.
+	stopRunning chan struct{}
+	// cancelDialings is used by Close() to stop any in-progress dialing.
+	cancelDialings context.CancelFunc
+	pending        chan *msgToSend
+	wg             sync.WaitGroup
 
 	muconn sync.Mutex
 	conn   net.Conn
 }
 
+type dialer interface {
+	DialContext(ctx context.Context, network, address string) (net.Conn, error)
+}
+
 // New creates a new Logger.
 func New(config Config) (*Fluent, error) {
 	if config.Timeout == 0 {
@@ -104,10 +119,6 @@ func New(config Config) (*Fluent, error) {
 	})
 }
 
-type dialer interface {
-	Dial(string, string) (net.Conn, error)
-}
-
 func newWithDialer(config Config, d dialer) (f *Fluent, err error) {
 	if config.FluentNetwork == "" {
 		config.FluentNetwork = defaultNetwork
@@ -142,20 +153,24 @@ func newWithDialer(config Config, d dialer) (f *Fluent, err error) {
 	}
 
 	if config.Async {
+		ctx, cancel := context.WithCancel(context.Background())
+
 		f = &Fluent{
-			Config:  config,
-			dialer:  d,
-			pending: make(chan *msgToSend, config.BufferLimit),
-			stopRunning: make(chan bool, 1),
+			Config:         config,
+			dialer:         d,
+			stopRunning:    make(chan struct{}),
+			cancelDialings: cancel,
+			pending:        make(chan *msgToSend, config.BufferLimit),
 		}
+
 		f.wg.Add(1)
-		go f.run()
+		go f.run(ctx)
 	} else {
 		f = &Fluent{
 			Config: config,
 			dialer: d,
 		}
-		err = f.connect()
+		err = f.connect(context.Background())
 	}
 	return
 }
@@ -252,7 +267,7 @@ func (f *Fluent) postRawData(msg *msgToSend) error {
 		return f.appendBuffer(msg)
 	}
 	// Synchronous write
-	return f.write(msg)
+	return f.write(context.Background(), msg)
 }
 
 // For sending forward protocol adopted JSON
@@ -286,7 +301,7 @@ func getUniqueID(timeUnix int64) (string, error) {
 		enc.Close()
 		return "", err
 	}
-	if err := binary.Write(enc, binary.LittleEndian, rand.Uint64()); err != nil {
+	if err := binary.Write(enc, binary.LittleEndian, randomGenerator()); err != nil {
 		enc.Close()
 		return "", err
 	}
@@ -326,13 +341,18 @@ func (f *Fluent) EncodeData(tag string, tm time.Time, message interface{}) (msg
 func (f *Fluent) Close() (err error) {
 	if f.Config.Async {
 		if f.Config.ForceStopAsyncSend {
-			f.stopRunning <- true
 			close(f.stopRunning)
+			f.cancelDialings()
 		}
+
 		close(f.pending)
 		f.wg.Wait()
 	}
-	f.close(f.conn)
+
+	f.muconn.Lock()
+	f.close()
+	f.muconn.Unlock()
+
 	return
 }
 
@@ -346,58 +366,91 @@ func (f *Fluent) appendBuffer(msg *msgToSend) error {
 	return nil
 }
 
-// close closes the connection.
-func (f *Fluent) close(c net.Conn) {
-	f.muconn.Lock()
-	if f.conn != nil && f.conn == c {
+// close closes the connection. Callers should take care of locking muconn first.
+func (f *Fluent) close() {
+	if f.conn != nil {
 		f.conn.Close()
 		f.conn = nil
 	}
-	f.muconn.Unlock()
 }
 
 // connect establishes a new connection using the specified transport.
-func (f *Fluent) connect() (err error) {
+func (f *Fluent) connect(ctx context.Context) (err error) {
 	switch f.Config.FluentNetwork {
 	case "tcp":
-		f.conn, err = f.dialer.Dial(
+		f.conn, err = f.dialer.DialContext(ctx,
 			f.Config.FluentNetwork,
 			f.Config.FluentHost+":"+strconv.Itoa(f.Config.FluentPort))
 	case "unix":
-		f.conn, err = f.dialer.Dial(
+		f.conn, err = f.dialer.DialContext(ctx,
 			f.Config.FluentNetwork,
 			f.Config.FluentSocketPath)
 	default:
 		err = NewErrUnknownNetwork(f.Config.FluentNetwork)
 	}
+
 	return err
 }
 
-func (f *Fluent) run() {
-	drainEvents := false
-	var emitEventDrainMsg sync.Once
+var errIsClosing = errors.New("fluent logger is closing")
+
+func (f *Fluent) connectOrRetry(ctx context.Context) error {
+	// Use a Time channel instead of time.Sleep() to avoid blocking this
+	// goroutine during possibly way too much time (because of the exponential
+	// back-off retry).
+	waiter := time.After(time.Duration(0))
+	for i := 0; i < f.Config.MaxRetry; i++ {
+		select {
+		case <-waiter:
+			err := f.connect(ctx)
+			if err == nil {
+				return nil
+			}
+
+			if _, ok := err.(*ErrUnknownNetwork); ok {
+				return err
+			}
+			if err == context.Canceled {
+				return errIsClosing
+			}
+
+			waitTime := f.Config.RetryWait * e(defaultReconnectWaitIncreRate, float64(i-1))
+			if waitTime > f.Config.MaxRetryWait {
+				waitTime = f.Config.MaxRetryWait
+			}
+
+			waiter = time.After(time.Duration(waitTime) * time.Millisecond)
+
+			fmt.Fprintf(os.Stderr, "[%s] An error happened during connect: %s. Retrying to connect in %dms.", time.Now().Format(time.RFC3339), err, waitTime)
+		case <-ctx.Done():
+			return errIsClosing
+		}
+	}
+
+	return fmt.Errorf("could not connect to fluentd after %d retries", f.Config.MaxRetry)
+}
+
+// run is the goroutine used to unqueue and write logs in async mode. That
+// goroutine is meant to run during the whole life of the Fluent logger.
+func (f *Fluent) run(ctx context.Context) {
 	for {
 		select {
 		case entry, ok := <-f.pending:
+			// f.stopRunning is closed before f.pending only when ForceStopAsyncSend
+			// is enabled. Otherwise, f.pending is closed when Close() is called.
 			if !ok {
 				f.wg.Done()
 				return
 			}
-			if drainEvents {
-				emitEventDrainMsg.Do(func() { fmt.Fprintf(os.Stderr, "[%s] Discarding queued events...\n", time.Now().Format(time.RFC3339)) })
-				continue
-			}
-			err := f.write(entry)
-			if err != nil {
+
+			if err := f.write(ctx, entry); err != nil && err != errIsClosing {
 				fmt.Fprintf(os.Stderr, "[%s] Unable to send logs to fluentd, reconnecting...\n", time.Now().Format(time.RFC3339))
 			}
-		}
-		select {
-		case stopRunning, ok := <-f.stopRunning:
-			if stopRunning || !ok {
-				drainEvents = true
-			}
-		default:
+		case <-f.stopRunning:
+			fmt.Fprintf(os.Stderr, "[%s] Discarding queued events...\n", time.Now().Format(time.RFC3339))
+
+			f.wg.Done()
+			return
 		}
 	}
 }
@@ -406,63 +459,59 @@ func e(x, y float64) int {
 	return int(math.Pow(x, y))
 }
 
-func (f *Fluent) write(msg *msgToSend) error {
-	var c net.Conn
-	for i := 0; i < f.Config.MaxRetry; i++ {
-		c = f.conn
-		// Connect if needed
-		if c == nil {
-			f.muconn.Lock()
-			if f.conn == nil {
-				err := f.connect()
-				if err != nil {
-					f.muconn.Unlock()
-
-					if _, ok := err.(*ErrUnknownNetwork); ok {
-						// do not retry on unknown network error
-						break
-					}
-					waitTime := f.Config.RetryWait * e(defaultReconnectWaitIncreRate, float64(i-1))
-					if waitTime > f.Config.MaxRetryWait {
-						waitTime = f.Config.MaxRetryWait
-					}
-					time.Sleep(time.Duration(waitTime) * time.Millisecond)
-					continue
-				}
+func (f *Fluent) write(ctx context.Context, msg *msgToSend) error {
+	writer := func() (bool, error) {
+		// This function is used to ensure muconn is properly locked and unlocked
+		// between each retry. This gives the importunity to other goroutines to
+		// lock it (e.g. to close the connection).
+		f.muconn.Lock()
+		defer f.muconn.Unlock()
+
+		if f.conn == nil {
+			if err := f.connectOrRetry(ctx); err != nil {
+				return false, err
 			}
-			c = f.conn
-			f.muconn.Unlock()
 		}
 
-		// We're connected, write msg
 		t := f.Config.WriteTimeout
 		if time.Duration(0) < t {
-			c.SetWriteDeadline(time.Now().Add(t))
+			f.conn.SetWriteDeadline(time.Now().Add(t))
 		} else {
-			c.SetWriteDeadline(time.Time{})
+			f.conn.SetWriteDeadline(time.Time{})
 		}
-		_, err := c.Write(msg.data)
+
+		_, err := f.conn.Write(msg.data)
 		if err != nil {
-			f.close(c)
-		} else {
-			// Acknowledgment check
-			if msg.ack != "" {
-				resp := &AckResp{}
-				if f.Config.MarshalAsJSON {
-					dec := json.NewDecoder(c)
-					err = dec.Decode(resp)
-				} else {
-					r := msgp.NewReader(c)
-					err = resp.DecodeMsg(r)
-				}
-				if err != nil || resp.Ack != msg.ack {
-					f.close(c)
-					continue
-				}
+			f.close()
+			return true, err
+		}
+
+		// Acknowledgment check
+		if msg.ack != "" {
+			resp := &AckResp{}
+			if f.Config.MarshalAsJSON {
+				dec := json.NewDecoder(f.conn)
+				err = dec.Decode(resp)
+			} else {
+				r := msgp.NewReader(f.conn)
+				err = resp.DecodeMsg(r)
 			}
+
+			if err != nil || resp.Ack != msg.ack {
+				fmt.Fprintf(os.Stderr, "fluent#write: message ack (%s) doesn't match expected one (%s). Closing connection...", resp.Ack, msg.ack)
+				f.close()
+				return true, err
+			}
+		}
+
+		return false, nil
+	}
+
+	for i := 0; i < f.Config.MaxRetry; i++ {
+		if retry, err := writer(); !retry {
 			return err
 		}
 	}
 
-	return fmt.Errorf("fluent#write: failed to reconnect, max retry: %v", f.Config.MaxRetry)
+	return fmt.Errorf("fluent#write: failed to write after %d attempts", f.Config.MaxRetry)
 }