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() and a new stopAsyncConnect channel has been introduced to stop it. 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. 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 <albin.kerouanton@knplabs.com>

Author: Albin Kerouanton

fluent/fluent.go

@@ -1,6 +1,7 @@
 package fluent
 
 import (
+"context"
 "encoding/json"
 "errors"
 "fmt"
@@ -85,15 +86,24 @@ type msgToSend struct {
 type Fluent struct {
 Config
 
-dialer dialer
+dialer dialer
+// stopRunning is used in async mode to signal to run() and connectOrRetryAsync()
+// they should abort.
 stopRunning chan bool
-pending chan *msgToSend
-wg sync.WaitGroup
+// stopAsyncConnect is used by connectOrRetryAsync() to signal to
+// connectOrRetry() it should abort.
+stopAsyncConnect chan bool
+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 +114,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
@@ -143,9 +149,11 @@ func newWithDialer(config Config, d dialer) (f *Fluent, err error) {
 
 if config.Async {
 f = &Fluent{
-Config: config,
-dialer: d,
-pending: make(chan *msgToSend, config.BufferLimit),
+Config: config,
+dialer: d,
+stopRunning: make(chan bool),
+stopAsyncConnect: make(chan bool),
+pending: make(chan *msgToSend, config.BufferLimit),
 }
 f.wg.Add(1)
 go f.run()
@@ -154,7 +162,7 @@ func newWithDialer(config Config, d dialer) (f *Fluent, err error) {
 Config: config,
 dialer: d,
 }
-err = f.connect()
+err = f.connect(context.Background())
 }
 return
 }
@@ -331,7 +339,11 @@ func (f *Fluent) Close() (err error) {
 close(f.pending)
 f.wg.Wait()
 }
-f.close(f.conn)
+
+f.muconn.Lock()
+f.close()
+f.muconn.Unlock()
+
 return
 }
 
@@ -346,35 +358,105 @@ func (f *Fluent) appendBuffer(msg *msgToSend) error {
 }
 
 // close closes the connection.
-func (f *Fluent) close(c net.Conn) {
-f.muconn.Lock()
-if f.conn != nil && f.conn == c {
+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
 }
 
+var errIsClosing = errors.New("fluent logger is closing, aborting async connect")
+
+func (f *Fluent) connectOrRetry(ctx context.Context) error {
+// Use a Time channel instead of time.Sleep() to avoid blocking this
+// goroutine during eventually 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 {
+// No need to retry on unknown network error. Thus false is passed
+// to ready channel to let the other end drain the message queue.
+return err
+}
+
+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 <-f.stopAsyncConnect:
+return errIsClosing
+}
+}
+
+return fmt.Errorf("could not connect to fluentd after %d retries", f.Config.MaxRetry)
+}
+
+// connectOrRetryAsync returns an error when it fails to connect to fluentd or
+// when Close() has been called.
+func (f *Fluent) connectOrRetryAsync(ctx context.Context) error {
+ctx, cancelDialing := context.WithCancel(ctx)
+errCh := make(chan error)
+
+f.wg.Add(1)
+go func(ctx context.Context, errCh chan<- error) {
+defer f.wg.Done()
+errCh <- f.connectOrRetry(ctx)
+}(ctx, errCh)
+
+for {
+select {
+case _, ok := <-f.stopRunning:
+// If f.stopRunning is closed before we got something on errCh,
+// we need to wait a bit more.
+if !ok {
+break
+}
+
+// Stop any connection dialing and then tell connectOrRetry to stop
+// trying to dial the connection. This has to be done in this
+// specifc order to make sure connectOrRetry() is not blocking on
+// the connection dialing.
+cancelDialing()
+f.stopAsyncConnect <- true
+case err := <-errCh:
+return err
+}
+}
+}
+
+// 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() {
 drainEvents := false
 var emitEventDrainMsg sync.Once
+
 for {
 select {
 case entry, ok := <-f.pending:
@@ -387,16 +469,16 @@ func (f *Fluent) run() {
 continue
 }
 err := f.write(entry)
-if err != nil {
+if err == errIsClosing {
+drainEvents = true
+} else if err != nil {
+// TODO: log failing message?
 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:
 }
 }
 }
@@ -406,62 +488,64 @@ func e(x, y float64) int {
 }
 
 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
-}
+// 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).
+writer := func() (bool, error) {
+f.muconn.Lock()
+defer f.muconn.Unlock()
+
+if f.conn == nil {
+var err error
+if f.Config.Async {
+err = f.connectOrRetryAsync(context.Background())
+} else {
+err = f.connectOrRetry(context.Background())
+}
+
+if 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 {
+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)
 }

fluent/fluent_test.go

@@ -2,6 +2,7 @@ package fluent
 
 import (
 "bytes"
+"context"
 "encoding/json"
 "errors"
 "io/ioutil"
@@ -29,9 +30,14 @@ type testDialer struct {
 connCh chan *Conn
 }
 
-func (d *testDialer) Dial(string, string) (net.Conn, error) {
-conn := <-d.connCh
-if conn == nil {
+func (d *testDialer) DialContext(ctx context.Context, _, _ string) (net.Conn, error) {
+select {
+case conn := <-d.connCh:
+if conn == nil {
+return nil, errors.New("failed to dial")
+}
+return conn, nil
+case <-ctx.Done():
 return nil, errors.New("failed to dial")
 }
 return conn, nil
@@ -404,8 +410,6 @@ func timeout(t *testing.T, duration time.Duration, fn func(), reason string) {
 }
 
 func TestCloseOnFailingAsyncConnect(t *testing.T) {
-t.Skip("Broken tests")
-
 testcases := map[string]Config{
 "with ForceStopAsyncSend and with RequestAck": {
 Async: true,
@@ -457,8 +461,6 @@ func ackRespMsgp(t *testing.T, ack string) string {
 }
 
 func TestCloseOnFailingAsyncReconnect(t *testing.T) {
-t.Skip("Broken tests")
-
 testcases := map[string]Config{
 "with RequestAck": {
 Async: true,