[ADDED] Automatic reconnect on write error option (#2055)

Signed-off-by: Piotr Piotrowski <piotr@synadia.com>

Author: piotrpio

nats.go

@@ -410,6 +410,34 @@ type Options struct {
 	// Defaults to 1m.
 	FlusherTimeout time.Duration
 
+	// ReconnectOnFlusherError, when set to true, causes the client to
+	// trigger a reconnect if the background flusher fails to write to the
+	// underlying connection for any reason (timeout, broken pipe,
+	// connection reset, EOF etc.).
+	//
+	// This is an advanced option. Most applications do not need to enable
+	// it: the server-side stale connection detection (via PingInterval /
+	// MaxPingsOut) and the read loop's own error handling will eventually
+	// notice a dead connection and the client will reconnect. Enable this
+	// only if you need faster recovery from a stalled or broken TCP write
+	// — for example, in latency-sensitive setups where waiting for a ping
+	// timeout is unacceptable.
+	//
+	// Messages buffered at the time of the error are lost, as they are
+	// with any flusher write error. The purpose of this option is to
+	// limit the blast radius by preventing further messages from being
+	// buffered into a potentially corrupted connection, not to recover
+	// the in-flight data.
+	//
+	// When triggered, the standard DisconnectErrHandler and
+	// ReconnectHandler callbacks are invoked as with any other reconnect.
+	// The first reconnect attempt bypasses the configured ReconnectWait
+	// so that recovery is as fast as possible; if that attempt fails,
+	// subsequent attempts obey the normal backoff.
+	//
+	// Defaults to false.
+	ReconnectOnFlusherError bool
+
 	// PingInterval is the period at which the client will be sending ping
 	// commands to the server, disabled if 0 or negative.
 	// Defaults to 2m.
@@ -1211,6 +1239,17 @@ func FlusherTimeout(t time.Duration) Option {
 	}
 }
 
+// ReconnectOnFlusherError is an Option to automatically trigger a
+// reconnect when the background flusher hits any write error. See
+// [Options.ReconnectOnFlusherError] for details. This is an
+// advanced option and is usually not required.
+func ReconnectOnFlusherError() Option {
+	return func(o *Options) error {
+		o.ReconnectOnFlusherError = true
+		return nil
+	}
+}
+
 // DrainTimeout is an Option to set the timeout for draining a connection.
 // Defaults to 30s.
 func DrainTimeout(t time.Duration) Option {
@@ -2440,8 +2479,10 @@ func (nc *Conn) ForceReconnect() error {
 	// Stop ping timer if set.
 	nc.stopPingTimer()
 
-	// Go ahead and make sure we have flushed the outbound
+	// flush any pending data and switch to pending mode to buffer new outgoing
+	// data until we reconnect and can flush it.
 	nc.bw.flush()
+	nc.bw.switchToPending()
 	nc.conn.Close()
 
 	nc.changeConnStatus(RECONNECTING)
@@ -3342,8 +3383,10 @@ func (nc *Conn) doReconnect(err error, forceReconnect bool) {
 }
 
 // processOpErr handles errors from reading or parsing the protocol.
-// The lock should not be held entering this function.
-func (nc *Conn) processOpErr(err error) bool {
+// The lock should not be held entering this function. If forceReconnect
+// is true, the first reconnect attempt will bypass the configured
+// ReconnectWait; subsequent attempts still obey the normal backoff.
+func (nc *Conn) processOpErr(err error, forceReconnect bool) bool {
 	nc.mu.Lock()
 	defer nc.mu.Unlock()
 	if nc.isConnecting() || nc.isClosed() || nc.isReconnecting() {
@@ -3366,7 +3409,7 @@ func (nc *Conn) processOpErr(err error) bool {
 		// Clear any queued pongs, e.g. pending flush calls.
 		nc.clearPendingFlushCalls()
 
-		go nc.doReconnect(err, false)
+		go nc.doReconnect(err, forceReconnect)
 		return false
 	}
 
@@ -3469,7 +3512,7 @@ func (nc *Conn) readLoop() {
 			err = nc.parse(buf)
 		}
 		if err != nil {
-			if shouldClose := nc.processOpErr(err); shouldClose {
+			if shouldClose := nc.processOpErr(err, false); shouldClose {
 				nc.close(CLOSED, true, nil)
 			}
 			break
@@ -3917,6 +3960,13 @@ func (nc *Conn) flusher() {
 				if asyncErrorCB := nc.Opts.AsyncErrorCB; asyncErrorCB != nil {
 					nc.ach.push(func() { asyncErrorCB(nc, nil, err) })
 				}
+				if nc.Opts.ReconnectOnFlusherError {
+					nc.mu.Unlock()
+					if shouldClose := nc.processOpErr(err, true); shouldClose {
+						nc.close(CLOSED, true, nil)
+					}
+					return
+				}
 			}
 		}
 		nc.mu.Unlock()
@@ -4096,11 +4146,11 @@ func (nc *Conn) processErr(ie string) {
 
 	// FIXME(dlc) - process Slow Consumer signals special.
 	if e == STALE_CONNECTION {
-		close = nc.processOpErr(ErrStaleConnection)
+		close = nc.processOpErr(ErrStaleConnection, false)
 	} else if e == MAX_CONNECTIONS_ERR {
-		close = nc.processOpErr(ErrMaxConnectionsExceeded)
+		close = nc.processOpErr(ErrMaxConnectionsExceeded, false)
 	} else if e == MAX_ACCOUNT_CONNECTIONS_ERR {
-		close = nc.processOpErr(ErrMaxAccountConnectionsExceeded)
+		close = nc.processOpErr(ErrMaxAccountConnectionsExceeded, false)
 	} else if strings.HasPrefix(e, PERMISSIONS_ERR) {
 		nc.processTransientError(fmt.Errorf("%w: %s", ErrPermissionViolation, ne))
 	} else if strings.HasPrefix(e, MAX_SUBSCRIPTIONS_ERR) {
@@ -5682,7 +5732,7 @@ func (nc *Conn) processPingTimer() {
 	nc.pout++
 	if nc.pout > nc.Opts.MaxPingsOut {
 		nc.mu.Unlock()
-		if shouldClose := nc.processOpErr(ErrStaleConnection); shouldClose {
+		if shouldClose := nc.processOpErr(ErrStaleConnection, false); shouldClose {
 			nc.close(CLOSED, true, nil)
 		}
 		return

test/conn_test.go

@@ -1931,7 +1931,7 @@ func TestCustomFlusherTimeout(t *testing.T) {
 
 	errCh := make(chan error, 1)
 	wg := sync.WaitGroup{}
-	wg.Add(2)
+	wg.Add(1)
 	go func() {
 		defer wg.Done()
 		for {
@@ -1949,50 +1949,7 @@ func TestCustomFlusherTimeout(t *testing.T) {
 	}()
 	defer nc1.Close()
 
-	l, e := net.Listen("tcp", "127.0.0.1:0")
-	if e != nil {
-		t.Fatal("Could not listen on an ephemeral port")
-	}
-	tl := l.(*net.TCPListener)
-	defer tl.Close()
-
-	addr := tl.Addr().(*net.TCPAddr)
-
-	fsDoneCh := make(chan struct{}, 1)
-	fsErrCh := make(chan error, 1)
-	go func() {
-		defer wg.Done()
-
-		serverInfo := "INFO {\"server_id\":\"foobar\",\"host\":\"%s\",\"port\":%d,\"auth_required\":false,\"tls_required\":false,\"max_payload\":%d}\r\n"
-		conn, err := l.Accept()
-		if err != nil {
-			fsErrCh <- err
-			return
-		}
-		defer conn.Close()
-		// Make it small on purpose
-		if err := conn.(*net.TCPConn).SetReadBuffer(1024); err != nil {
-			fsErrCh <- err
-			return
-		}
-
-		info := fmt.Sprintf(serverInfo, addr.IP, addr.Port, 1024*1024)
-		conn.Write([]byte(info))
-
-		// Read connect and ping commands sent from the client
-		line := make([]byte, 100)
-		_, err = conn.Read(line)
-		if err != nil {
-			fsErrCh <- fmt.Errorf("Expected CONNECT and PING from client, got: %v", err)
-			return
-		}
-		conn.Write([]byte("PONG\r\n"))
-
-		// Don't consume anything at this point and wait to be notified
-		// that we are done.
-		<-fsDoneCh
-		fsErrCh <- nil
-	}()
+	addr := startStalledMockServer(t)
 
 	nc2, err := nats.Connect(
 		// URL to fake server
@@ -2046,17 +2003,12 @@ forLoop:
 		}
 	}
 
-	// Notify fake server that it can stop
-	close(fsDoneCh)
-
-	// Wait for go routines to end
+	// Close nc2 to fire its ClosedHandler, which signals the publisher
+	// goroutine (via doneCh) to exit. The fake server is torn down by
+	// the t.Cleanup registered in startStalledFakeServer.
+	nc2.Close()
 	wg.Wait()
 
-	// Make sure there were no error in the fake server
-	if err := <-fsErrCh; err != nil {
-		t.Fatalf("Fake server reported: %v", err)
-	}
-
 	// One of those two are guaranteed to be set.
 	err = nc2Err
 	if err == nil {
@@ -2082,6 +2034,178 @@ forLoop:
 	}
 }
 
+// startStalledMockServer starts a fake NATS server on an ephemeral port
+// that completes the INFO/CONNECT/PING handshake and then stops reading
+// from the socket, so client writes eventually stall.
+func startStalledMockServer(t *testing.T) *net.TCPAddr {
+	t.Helper()
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("Could not listen on an ephemeral port: %v", err)
+	}
+	addr := l.Addr().(*net.TCPAddr)
+
+	done := make(chan struct{})
+	exited := make(chan struct{})
+
+	go func() {
+		defer close(exited)
+		conn, err := l.Accept()
+		if err != nil {
+			return
+		}
+		defer conn.Close()
+		if err := conn.(*net.TCPConn).SetReadBuffer(1024); err != nil {
+			t.Errorf("Expected SetReadBuffer to succeed, got: %v", err)
+			return
+		}
+		info := fmt.Sprintf(
+			"INFO {\"server_id\":\"foobar\",\"host\":\"%s\",\"port\":%d,\"auth_required\":false,\"tls_required\":false,\"max_payload\":%d}\r\n",
+			addr.IP, addr.Port, 1024*1024,
+		)
+		conn.Write([]byte(info))
+		line := make([]byte, 100)
+		if _, err := conn.Read(line); err != nil {
+			t.Errorf("Expected CONNECT+PING, got: %v", err)
+			return
+		}
+		conn.Write([]byte("PONG\r\n"))
+		<-done
+	}()
+
+	t.Cleanup(func() {
+		close(done)
+		l.Close()
+		<-exited
+	})
+	return addr
+}
+
+func TestReconnectOnFlusherError(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.SkipNow()
+	}
+
+	for _, tc := range []struct {
+		name          string
+		withOption    bool
+		noReconnect   bool
+		wantReconnect bool
+		wantClosed    bool
+	}{
+		{"enabled", true, false, true, false},
+		{"disabled", false, false, false, false},
+		{"no_reconnect", true, true, false, true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			s := RunDefaultServer()
+			defer s.Shutdown()
+
+			fakeAddr := startStalledMockServer(t)
+
+			reconnectedCh := make(chan struct{}, 1)
+			closedCh := make(chan struct{}, 1)
+			asyncErrCh := make(chan error, 1)
+
+			opts := []nats.Option{
+				nats.SetCustomDialer(&lowWriteBufferDialer{}),
+				nats.FlusherTimeout(15 * time.Millisecond),
+				nats.MaxReconnects(10),
+				nats.DontRandomize(),
+				nats.ReconnectHandler(func(_ *nats.Conn) {
+					select {
+					case reconnectedCh <- struct{}{}:
+					default:
+					}
+				}),
+				nats.ClosedHandler(func(_ *nats.Conn) {
+					select {
+					case closedCh <- struct{}{}:
+					default:
+					}
+				}),
+				nats.ErrorHandler(func(_ *nats.Conn, _ *nats.Subscription, err error) {
+					select {
+					case asyncErrCh <- err:
+					default:
+					}
+				}),
+			}
+			if tc.withOption {
+				opts = append(opts, nats.ReconnectOnFlusherError())
+			}
+			if tc.noReconnect {
+				opts = append(opts, nats.NoReconnect())
+			}
+
+			nc, err := nats.Connect(
+				fmt.Sprintf("nats://127.0.0.1:%d,%s", fakeAddr.Port, nats.DefaultURL),
+				opts...,
+			)
+			if err != nil {
+				t.Fatalf("Connect: %v", err)
+			}
+			defer nc.Close()
+
+			if url := nc.ConnectedUrl(); url != fmt.Sprintf("nats://127.0.0.1:%d", fakeAddr.Port) {
+				t.Fatalf("Expected initial connection to fake server, got %q", url)
+			}
+
+			stopPub := make(chan struct{})
+			pubDone := make(chan struct{})
+			go func() {
+				defer close(pubDone)
+				tick := time.NewTicker(50 * time.Millisecond)
+				defer tick.Stop()
+				// Slightly under the library's internal buffer size so the
+				// flusher goroutine (not the Publish call) triggers the write.
+				payload := make([]byte, 32*1024-200)
+				for {
+					select {
+					case <-stopPub:
+						return
+					case <-tick.C:
+						nc.Publish("hello", payload)
+					}
+				}
+			}()
+			defer func() {
+				close(stopPub)
+				<-pubDone
+			}()
+
+			// Confirm the flusher actually observes a write error.
+			select {
+			case <-asyncErrCh:
+			case <-time.After(3 * time.Second):
+				t.Fatal("flusher did not report an async write error within 3s")
+			}
+
+			if tc.wantReconnect {
+				select {
+				case <-reconnectedCh:
+				case <-time.After(5 * time.Second):
+					t.Fatal("expected reconnect after flusher error")
+				}
+				if url := nc.ConnectedUrl(); url != s.ClientURL() {
+					t.Fatalf("expected to be reconnected to real server %q, got %q", s.ClientURL(), url)
+				}
+			} else if tc.wantClosed {
+				WaitOnChannel(t, closedCh, struct{}{})
+				if !nc.IsClosed() {
+					t.Fatalf("expected IsClosed() to be true, got status %v", nc.Status())
+				}
+			} else {
+				select {
+				case <-reconnectedCh:
+					t.Fatal("unexpected reconnect when ReconnectOnFlusherError is disabled")
+				case <-time.After(500 * time.Millisecond):
+				}
+			}
+		})
+	}
+}
+
 func TestNewServers(t *testing.T) {
 	s1Opts := test.DefaultTestOptions
 	s1Opts.Host = "127.0.0.1"