GODRIVER-2577 Retry heartbeat on timeout to prevent pool cleanup in F…

…AAS pause. (#1133)

x/mongo/driver/topology/server.go

@@ -520,6 +520,7 @@ func (s *Server) update() {
 		}
 	}
 
+	timeoutCnt := 0
 	for {
 		// Check if the server is disconnecting. Even if waitForNextCheck has already read from the done channel, we
 		// can safely read from it again because Disconnect closes the channel.
@@ -545,18 +546,42 @@ func (s *Server) update() {
 			continue
 		}
 
-		// Must hold the processErrorLock while updating the server description and clearing the
-		// pool. Not holding the lock leads to possible out-of-order processing of pool.clear() and
-		// pool.ready() calls from concurrent server description updates.
-		s.processErrorLock.Lock()
-		s.updateDescription(desc)
-		if err := desc.LastError; err != nil {
-			// Clear the pool once the description has been updated to Unknown. Pass in a nil service ID to clear
-			// because the monitoring routine only runs for non-load balanced deployments in which servers don't return
-			// IDs.
-			s.pool.clear(err, nil)
+		if isShortcut := func() bool {
+			// Must hold the processErrorLock while updating the server description and clearing the
+			// pool. Not holding the lock leads to possible out-of-order processing of pool.clear() and
+			// pool.ready() calls from concurrent server description updates.
+			s.processErrorLock.Lock()
+			defer s.processErrorLock.Unlock()
+
+			s.updateDescription(desc)
+			// Retry after the first timeout before clearing the pool in case of a FAAS pause as
+			// described in GODRIVER-2577.
+			if err := unwrapConnectionError(desc.LastError); err != nil && timeoutCnt < 1 {
+				if err == context.Canceled || err == context.DeadlineExceeded {
+					timeoutCnt++
+					// We want to immediately retry on timeout error. Continue to next loop.
+					return true
+				}
+				if err, ok := err.(net.Error); ok && err.Timeout() {
+					timeoutCnt++
+					// We want to immediately retry on timeout error. Continue to next loop.
+					return true
+				}
+			}
+			if err := desc.LastError; err != nil {
+				// Clear the pool once the description has been updated to Unknown. Pass in a nil service ID to clear
+				// because the monitoring routine only runs for non-load balanced deployments in which servers don't return
+				// IDs.
+				s.pool.clear(err, nil)
+			}
+			// We're either not handling a timeout error, or we just handled the 2nd consecutive
+			// timeout error. In either case, reset the timeout count to 0 and return false to
+			// continue the normal check process.
+			timeoutCnt = 0
+			return false
+		}(); isShortcut {
+			continue
 		}
-		s.processErrorLock.Unlock()
 
 		// If the server supports streaming or we're already streaming, we want to move to streaming the next response
 		// without waiting. If the server has transitioned to Unknown from a network error, we want to do another
@@ -707,19 +732,31 @@ func (s *Server) check() (description.Server, error) {
 	var err error
 	var durationNanos int64
 
-	// Create a new connection if this is the first check, the connection was closed after an error during the previous
-	// check, or the previous check was cancelled.
+	start := time.Now()
 	if s.conn == nil || s.conn.closed() || s.checkWasCancelled() {
+		// Create a new connection if this is the first check, the connection was closed after an error during the previous
+		// check, or the previous check was cancelled.
+		isNilConn := s.conn == nil
+		if !isNilConn {
+			s.publishServerHeartbeatStartedEvent(s.conn.ID(), false)
+		}
 		// Create a new connection and add it's handshake RTT as a sample.
 		err = s.setupHeartbeatConnection()
+		durationNanos = time.Since(start).Nanoseconds()
 		if err == nil {
 			// Use the description from the connection handshake as the value for this check.
 			s.rttMonitor.addSample(s.conn.helloRTT)
 			descPtr = &s.conn.desc
+			if !isNilConn {
+				s.publishServerHeartbeatSucceededEvent(s.conn.ID(), durationNanos, s.conn.desc, false)
+			}
+		} else {
+			err = unwrapConnectionError(err)
+			if !isNilConn {
+				s.publishServerHeartbeatFailedEvent(s.conn.ID(), durationNanos, err, false)
+			}
 		}
-	}
-
-	if descPtr == nil && err == nil {
+	} else {
 		// An existing connection is being used. Use the server description properties to execute the right heartbeat.
 
 		// Wrap conn in a type that implements driver.StreamerConnection.
@@ -729,7 +766,6 @@ func (s *Server) check() (description.Server, error) {
 		streamable := previousDescription.TopologyVersion != nil
 
 		s.publishServerHeartbeatStartedEvent(s.conn.ID(), s.conn.getCurrentlyStreaming() || streamable)
-		start := time.Now()
 		switch {
 		case s.conn.getCurrentlyStreaming():
 			// The connection is already in a streaming state, so we stream the next response.

x/mongo/driver/topology/server_test.go

@@ -11,8 +11,12 @@ package topology
 
 import (
 	"context"
+	"crypto/tls"
+	"crypto/x509"
 	"errors"
+	"io/ioutil"
 	"net"
+	"os"
 	"runtime"
 	"sync"
 	"sync/atomic"
@@ -49,6 +53,144 @@ func (cncd *channelNetConnDialer) DialContext(_ context.Context, _, _ string) (n
 	return cnc, nil
 }
 
+type errorQueue struct {
+	errors []error
+	mutex  sync.Mutex
+}
+
+func (eq *errorQueue) head() error {
+	eq.mutex.Lock()
+	defer eq.mutex.Unlock()
+	if len(eq.errors) > 0 {
+		return eq.errors[0]
+	}
+	return nil
+}
+
+func (eq *errorQueue) dequeue() bool {
+	eq.mutex.Lock()
+	defer eq.mutex.Unlock()
+	if len(eq.errors) > 0 {
+		eq.errors = eq.errors[1:]
+		return true
+	}
+	return false
+}
+
+type timeoutConn struct {
+	net.Conn
+	errors *errorQueue
+}
+
+func (c *timeoutConn) Read(b []byte) (int, error) {
+	n, err := 0, c.errors.head()
+	if err == nil {
+		n, err = c.Conn.Read(b)
+	}
+	return n, err
+}
+
+func (c *timeoutConn) Write(b []byte) (int, error) {
+	n, err := 0, c.errors.head()
+	if err == nil {
+		n, err = c.Conn.Write(b)
+	}
+	return n, err
+}
+
+type timeoutDialer struct {
+	Dialer
+	errors *errorQueue
+}
+
+func (d *timeoutDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	c, e := d.Dialer.DialContext(ctx, network, address)
+
+	if caFile := os.Getenv("MONGO_GO_DRIVER_CA_FILE"); len(caFile) > 0 {
+		pem, err := ioutil.ReadFile(caFile)
+		if err != nil {
+			return nil, err
+		}
+
+		ca := x509.NewCertPool()
+		if !ca.AppendCertsFromPEM(pem) {
+			return nil, errors.New("unable to load CA file")
+		}
+
+		config := &tls.Config{
+			InsecureSkipVerify: true,
+			RootCAs:            ca,
+		}
+		c = tls.Client(c, config)
+	}
+	return &timeoutConn{c, d.errors}, e
+}
+
+// TestServerHeartbeatTimeout tests timeout retry for GODRIVER-2577.
+func TestServerHeartbeatTimeout(t *testing.T) {
+	networkTimeoutError := &net.DNSError{
+		IsTimeout: true,
+	}
+
+	testCases := []struct {
+		desc              string
+		ioErrors          []error
+		expectPoolCleared bool
+	}{
+		{
+			desc:              "one single timeout should not clear the pool",
+			ioErrors:          []error{nil, networkTimeoutError, nil, networkTimeoutError, nil},
+			expectPoolCleared: false,
+		},
+		{
+			desc:              "continuous timeouts should clear the pool",
+			ioErrors:          []error{nil, networkTimeoutError, networkTimeoutError, nil},
+			expectPoolCleared: true,
+		},
+	}
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.desc, func(t *testing.T) {
+			t.Parallel()
+
+			var wg sync.WaitGroup
+			wg.Add(1)
+
+			errors := &errorQueue{errors: tc.ioErrors}
+			tpm := monitor.NewTestPoolMonitor()
+			server := NewServer(
+				address.Address("localhost:27017"),
+				primitive.NewObjectID(),
+				WithConnectionPoolMonitor(func(*event.PoolMonitor) *event.PoolMonitor {
+					return tpm.PoolMonitor
+				}),
+				WithConnectionOptions(func(opts ...ConnectionOption) []ConnectionOption {
+					return append(opts,
+						WithDialer(func(d Dialer) Dialer {
+							var dialer net.Dialer
+							return &timeoutDialer{&dialer, errors}
+						}))
+				}),
+				WithServerMonitor(func(*event.ServerMonitor) *event.ServerMonitor {
+					return &event.ServerMonitor{
+						ServerHeartbeatStarted: func(e *event.ServerHeartbeatStartedEvent) {
+							if !errors.dequeue() {
+								wg.Done()
+							}
+						},
+					}
+				}),
+				WithHeartbeatInterval(func(time.Duration) time.Duration {
+					return 200 * time.Millisecond
+				}),
+			)
+			require.NoError(t, server.Connect(nil))
+			wg.Wait()
+			assert.Equal(t, tc.expectPoolCleared, tpm.IsPoolCleared(), "expected pool cleared to be %v but was %v", tc.expectPoolCleared, tpm.IsPoolCleared())
+		})
+	}
+}
+
 // TestServerConnectionTimeout tests how different timeout errors are handled during connection
 // creation and server handshake.
 func TestServerConnectionTimeout(t *testing.T) {