Try to asynchronously connect to Fluentd before writing

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 for tcp timeout.

To fix these edge cases, the connection isn't initialized lazily
anymore. However, it's still initialized asynchronously and with the
exponential back-off retry. The Fluent.run() method has been adapted to
wait for either the connection to become ready or to receive a stop
signal on the stopRunning channel before starting to unqueue logs.

Author: Albin Kerouanton

fluent/fluent.go

@@ -1,6 +1,7 @@
 package fluent
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -15,8 +16,9 @@ import (
 	"bytes"
 	"encoding/base64"
 	"encoding/binary"
-	"github.com/tinylib/msgp/msgp"
 	"math/rand"
+
+	"github.com/tinylib/msgp/msgp"
 )
 
 const (
@@ -84,6 +86,7 @@ type msgToSend struct {
 type Fluent struct {
 	Config
 
+	ready       chan bool
 	stopRunning chan bool
 	pending     chan *msgToSend
 	wg          sync.WaitGroup
@@ -130,14 +133,16 @@ func New(config Config) (f *Fluent, err error) {
 	}
 	if config.Async {
 		f = &Fluent{
-			Config:  config,
-			pending: make(chan *msgToSend, config.BufferLimit),
+			Config:      config,
+			ready:       make(chan bool),
+			stopRunning: make(chan bool),
+			pending:     make(chan *msgToSend, config.BufferLimit),
 		}
 		f.wg.Add(1)
 		go f.run()
 	} else {
 		f = &Fluent{Config: config}
-		err = f.connect()
+		err = f.connect(context.Background())
 	}
 	return
 }
@@ -339,22 +344,95 @@ func (f *Fluent) close(c net.Conn) {
 }
 
 // connect establishes a new connection using the specified transport.
-func (f *Fluent) connect() (err error) {
+func (f *Fluent) connect(ctx context.Context) (err error) {
+	f.muconn.Lock()
+	defer f.muconn.Unlock()
 
+	dialer := net.Dialer{Timeout: f.Config.Timeout}
 	switch f.Config.FluentNetwork {
 	case "tcp":
-		f.conn, err = net.DialTimeout(f.Config.FluentNetwork, f.Config.FluentHost+":"+strconv.Itoa(f.Config.FluentPort), f.Config.Timeout)
+		f.conn, err = dialer.DialContext(ctx,
+			f.Config.FluentNetwork,
+			f.Config.FluentHost+":"+strconv.Itoa(f.Config.FluentPort))
 	case "unix":
-		f.conn, err = net.DialTimeout(f.Config.FluentNetwork, f.Config.FluentSocketPath, f.Config.Timeout)
+		f.conn, err = dialer.DialContext(ctx,
+			f.Config.FluentNetwork,
+			f.Config.FluentSocketPath)
 	default:
 		err = NewErrUnknownNetwork(f.Config.FluentNetwork)
 	}
+
 	return err
 }
 
+func (f *Fluent) connectAsync(ctx context.Context, stopAsyncConnect <-chan bool) {
+	f.wg.Add(1)
+	defer f.wg.Done()
+
+	waiter := time.After(time.Duration(0))
+	for i := 0; i < f.Config.MaxRetry; i++ {
+		select {
+		case <-waiter:
+			err := f.connect(ctx)
+			if err == nil {
+				f.ready <- true
+				break
+			}
+
+			if _, ok := err.(*ErrUnknownNetwork); ok {
+				// No need to retry on unknown network error. Thus false is passed
+				// to ready channel to let the other end drain the message queue.
+				f.ready <- false
+				break
+			}
+
+			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)
+		case <-stopAsyncConnect:
+			break
+		}
+	}
+}
+
 func (f *Fluent) run() {
 	drainEvents := false
 	var emitEventDrainMsg sync.Once
+
+	// First we need to wait for the connection to become ready. We cannot
+	// initialize the connection lazily (eg. when the first message is
+	// received) because it'd be done in the first for-select iteration on
+	// f.pending and this would block the select without letting the chance to
+	// the select on f.stopRunning to signal its end to this goroutine.
+	ctx, cancelDialing := context.WithCancel(context.Background())
+	stopAsyncConnect := make(chan bool, 1)
+	go f.connectAsync(ctx, stopAsyncConnect)
+	for i := 0; i < f.Config.MaxRetry; i++ {
+		select {
+		case stopRunning, ok := <-f.stopRunning:
+			if stopRunning || !ok {
+				drainEvents = true
+			}
+			// Stop any connection dialing and then tell connectAsync to stop
+			// trying to dial the connection. This has to be done in this
+			// specifc order to make sure connectAsync() is not blocking on the
+			// connection dialing.
+			cancelDialing()
+			close(stopAsyncConnect)
+			break
+		case ready, ok := <-f.ready:
+			if !ready || !ok {
+				drainEvents = true
+			}
+			break
+		}
+	}
+
+	// At this point we can go ahead: the connection is either ready to use or
+	// drainEvents is true and thus all logs should be discarded.
 	for {
 		select {
 		case entry, ok := <-f.pending:
@@ -389,31 +467,6 @@ 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
-				}
-			}
-			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))