diff --git a/x/mongo/driver/topology/server.go b/x/mongo/driver/topology/server.go index 88b93b15e6..4a59643282 100644 --- a/x/mongo/driver/topology/server.go +++ b/x/mongo/driver/topology/server.go @@ -795,27 +795,30 @@ func (s *Server) check() (description.Server, error) { var duration time.Duration start := time.Now() + + // 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. 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. + connID := "0" if s.conn != nil { - s.publishServerHeartbeatStartedEvent(s.conn.ID(), false) + connID = s.conn.ID() } + s.publishServerHeartbeatStartedEvent(connID, false) // Create a new connection and add it's handshake RTT as a sample. err = s.setupHeartbeatConnection() duration = time.Since(start) + connID = "0" + if s.conn != nil { + connID = s.conn.ID() + } 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 s.conn != nil { - s.publishServerHeartbeatSucceededEvent(s.conn.ID(), duration, s.conn.desc, false) - } + s.publishServerHeartbeatSucceededEvent(connID, duration, s.conn.desc, false) } else { err = unwrapConnectionError(err) - if s.conn != nil { - s.publishServerHeartbeatFailedEvent(s.conn.ID(), duration, err, false) - } + s.publishServerHeartbeatFailedEvent(connID, duration, err, false) } } else { // An existing connection is being used. Use the server description properties to execute the right heartbeat. diff --git a/x/mongo/driver/topology/server_test.go b/x/mongo/driver/topology/server_test.go index ba92b6dd94..4a9ffb10a1 100644 --- a/x/mongo/driver/topology/server_test.go +++ b/x/mongo/driver/topology/server_test.go @@ -177,7 +177,12 @@ func TestServerHeartbeatTimeout(t *testing.T) { }), WithServerMonitor(func(*event.ServerMonitor) *event.ServerMonitor { return &event.ServerMonitor{ - ServerHeartbeatStarted: func(e *event.ServerHeartbeatStartedEvent) { + ServerHeartbeatSucceeded: func(e *event.ServerHeartbeatSucceededEvent) { + if !errors.dequeue() { + wg.Done() + } + }, + ServerHeartbeatFailed: func(e *event.ServerHeartbeatFailedEvent) { if !errors.dequeue() { wg.Done() }