GODRIVER-1708 Set starting time for monitoring connection handshake (#475)

Author: Divjot Arora

x/mongo/driver/topology/server.go

@@ -439,15 +439,17 @@ func (s *Server) updateDescription(desc description.Server) {
 }
 
 // createConnection creates a new connection instance but does not call connect on it. The caller must call connect
-// before the connection can be used for network operations.
-func (s *Server) createConnection(ctx context.Context) (*connection, error) {
+// before the connection can be used for network operations. The nowPtr pointer will be set to the time right before
+// the connection handshake starts and can be used by the caller to measure RTT for the connection handshake.
+func (s *Server) createConnection(ctx context.Context, nowPtr *time.Time) (*connection, error) {
 	opts := []ConnectionOption{
 		WithConnectTimeout(func(time.Duration) time.Duration { return s.cfg.heartbeatTimeout }),
 		WithReadTimeout(func(time.Duration) time.Duration { return s.cfg.heartbeatTimeout }),
 		WithWriteTimeout(func(time.Duration) time.Duration { return s.cfg.heartbeatTimeout }),
 		// We override whatever handshaker is currently attached to the options with a basic
 		// one because need to make sure we don't do auth.
 		WithHandshaker(func(h Handshaker) Handshaker {
+			*nowPtr = time.Now()
 			return operation.NewIsMaster().AppName(s.cfg.appname).Compressors(s.cfg.compressionOpts)
 		}),
 		// Override any command monitors specified in options with nil to avoid monitoring heartbeats.
@@ -487,7 +489,7 @@ func (s *Server) heartbeat(conn *connection) (description.Server, *connection) {
 		}
 
 		if conn == nil {
-			conn, err = s.createConnection(ctx)
+			conn, err = s.createConnection(ctx, &now)
 
 			conn.connect(ctx)
 

x/mongo/driver/topology/server_test.go

@@ -313,13 +313,34 @@ func TestServer(t *testing.T) {
 		)
 		assert.Nil(t, err, "NewServer error: %v", err)
 
-		conn, err := s.createConnection(context.Background())
+		var now time.Time
+		conn, err := s.createConnection(context.Background(), &now)
 		assert.Nil(t, err, "createConnection error: %v", err)
 
 		assert.Equal(t, s.cfg.heartbeatTimeout, 10*time.Second, "expected heartbeatTimeout to be: %v, got: %v", 10*time.Second, s.cfg.heartbeatTimeout)
 		assert.Equal(t, s.cfg.heartbeatTimeout, conn.readTimeout, "expected readTimeout to be: %v, got: %v", s.cfg.heartbeatTimeout, conn.readTimeout)
 		assert.Equal(t, s.cfg.heartbeatTimeout, conn.writeTimeout, "expected writeTimeout to be: %v, got: %v", s.cfg.heartbeatTimeout, conn.writeTimeout)
 	})
+	t.Run("heartbeat connection RTT is accurately measured", func(t *testing.T) {
+		dialer := &channelNetConnDialer{}
+
+		s, err := NewServer(
+			address.Address("localhost"),
+			WithConnectionOptions(func(connOpts ...ConnectionOption) []ConnectionOption {
+				return append(
+					connOpts,
+					WithDialer(func(Dialer) Dialer {
+						return dialer
+					}),
+				)
+			}),
+		)
+		assert.Nil(t, err, "NewServer error: %v", err)
+
+		_, _ = s.heartbeat(nil)
+		assert.True(t, s.averageRTTSet, "expected averageRTTSet to be true but was not")
+		assert.True(t, s.averageRTT < time.Second, "expected average RTT to be less than 1s, got %v", s.averageRTT)
+	})
 }
 
 func includesMetadata(t *testing.T, wm []byte) bool {