Replace pool.inUseConnections with server.operationCount and respond to PR feedback.

Author: matthewdale

mongo/integration/server_selection_prose_test.go

@@ -10,6 +10,7 @@ import (
 	"context"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -50,6 +51,13 @@ func TestServerSelectionProse(t *testing.T) {
 				AppName:         "loadBalancingTest",
 			},
 		})
+		// The automatic failpoint clearing may not clear failpoints set on specific hosts, so
+		// manually clear the failpoint we set on the specific mongos when the test is done.
+		defer func() {
+			mt.ResetClient(options.Client().
+				SetHosts([]string{failpointHost}))
+			mt.ClearFailPoints()
+		}()
 
 		// Reset the client with exactly 2 mongos hosts.
 		tpm := newTestPoolMonitor()
@@ -107,6 +115,11 @@ func TestServerSelectionProse(t *testing.T) {
 			SetHosts(hosts[:2]).
 			SetPoolMonitor(tpm.PoolMonitor))
 
+		// Sleep for 100ms to allow all server state discovery to complete. We need both servers to
+		// be selectable when we start running the test or the distribution of selected servers will
+		// be skewed. Unfortunately there's not currently another signal we can block on.
+		time.Sleep(100 * time.Millisecond)
+
 		// Start 25 goroutines that each run 10 findOne operations. Run 25 goroutines instead of the
 		// 10 that the prose test specifies to reduce intermittent test failures caused by the
 		// random selections not being perfectly even over small numbers of operations.

x/mongo/driver/topology/connection.go

@@ -593,6 +593,10 @@ type Connection struct {
 	refCount      int
 	cleanupPoolFn func()
 
+	// cleanupServerFn resets the server state when a connection is returned to the connection pool
+	// via Close() or expired via Expire().
+	cleanupServerFn func()
+
 	mu sync.RWMutex
 }
 
@@ -694,6 +698,10 @@ func (c *Connection) cleanupReferences() error {
 		c.cleanupPoolFn()
 		c.cleanupPoolFn = nil
 	}
+	if c.cleanupServerFn != nil {
+		c.cleanupServerFn()
+		c.cleanupServerFn = nil
+	}
 	c.connection = nil
 	return err
 }

x/mongo/driver/topology/pool.go

@@ -85,7 +85,6 @@ type pool struct {
 	nextID                       uint64 // nextID is the next pool ID for a new connection.
 	pinnedCursorConnections      uint64
 	pinnedTransactionConnections uint64
-	inUseConnections             int64
 
 	address       address.Address
 	minSize       uint64
@@ -452,8 +451,6 @@ func (p *pool) checkOut(ctx context.Context) (conn *connection, err error) {
 				ConnectionID: w.conn.poolID,
 			})
 		}
-
-		atomic.AddInt64(&p.inUseConnections, 1)
 		return w.conn, nil
 	}
 
@@ -483,8 +480,6 @@ func (p *pool) checkOut(ctx context.Context) (conn *connection, err error) {
 				ConnectionID: w.conn.poolID,
 			})
 		}
-
-		atomic.AddInt64(&p.inUseConnections, 1)
 		return w.conn, nil
 	case <-ctx.Done():
 		if p.monitor != nil {
@@ -588,13 +583,6 @@ func (p *pool) checkIn(conn *connection) error {
 		})
 	}
 
-	// Decrement the number of in-use connections once we know we have a non-nil connection that
-	// came from this pool. Do this in checkIn() instead of checkInNoEvent() because the latter is
-	// called by createConnections() and is not necessarily a check-in of an in-use connection. Use
-	// an int64 instead of a uint64 to mitigate the impact of any possible bugs that could cause the
-	// uint64 to underflow, which would effectively make the server unselectable.
-	atomic.AddInt64(&p.inUseConnections, -1)
-
 	return p.checkInNoEvent(conn)
 }
 
@@ -777,10 +765,6 @@ func (p *pool) availableConnectionCount() int {
 	return len(p.idleConns)
 }
 
-func (p *pool) inUseConnectionCount() int64 {
-	return atomic.LoadInt64(&p.inUseConnections)
-}
-
 // createConnections creates connections for wantConn requests on the newConnWait queue.
 func (p *pool) createConnections(ctx context.Context, wg *sync.WaitGroup) {
 	defer wg.Done()

x/mongo/driver/topology/server.go

@@ -82,11 +82,13 @@ func connectionStateString(state int64) string {
 
 // Server is a single server within a topology.
 type Server struct {
-	// connectionstate must be accessed using the atomic package and should be at the beginning of
-	// the struct.
+	// The following integer fields must be accessed using the atomic package and should be at the
+	// beginning of the struct.
 	// - atomic bug: https://pkg.go.dev/sync/atomic#pkg-note-BUG
 	// - suggested layout: https://go101.org/article/memory-layout.html
+
 	connectionstate int64
+	operationCount  int64
 
 	cfg     *serverConfig
 	address address.Address
@@ -267,12 +269,25 @@ func (s *Server) Connection(ctx context.Context) (driver.Connection, error) {
 		return nil, ErrServerClosed
 	}
 
-	connImpl, err := s.pool.checkOut(ctx)
+	// Increment the operation count before calling checkOut to make sure that all connection
+	// requests are included in the operation count, including those in the wait queue. If we got an
+	// error insted of a connection, immediately decrement the operation count.
+	atomic.AddInt64(&s.operationCount, 1)
+	conn, err := s.pool.checkOut(ctx)
 	if err != nil {
+		atomic.AddInt64(&s.operationCount, -1)
 		return nil, err
 	}
 
-	return &Connection{connection: connImpl}, nil
+	return &Connection{
+		connection: conn,
+		cleanupServerFn: func() {
+			// Decrement the operation count whenever the caller is done with the connection. Use an
+			// int64 instead of a uint64 to mitigate the impact of any possible bugs that could
+			// cause the uint64 to underflow, which would make the server much less selectable.
+			atomic.AddInt64(&s.operationCount, -1)
+		},
+	}, nil
 }
 
 // ProcessHandshakeError implements SDAM error handling for errors that occur before a connection
@@ -814,6 +829,11 @@ func (s *Server) MinRTT() time.Duration {
 	return s.rttMonitor.getMinRTT()
 }
 
+// OperationCount returns the current number of in-progress operations for this server.
+func (s *Server) OperationCount() int64 {
+	return atomic.LoadInt64(&s.operationCount)
+}
+
 // String implements the Stringer interface.
 func (s *Server) String() string {
 	desc := s.Description()

x/mongo/driver/topology/topology.go

@@ -440,7 +440,7 @@ func (t *Topology) SelectServer(ctx context.Context, ss description.ServerSelect
 		// Of the two randomly selected suitable servers, pick the one with fewer in-use connections.
 		// We use in-use connections as an analog for in-progress operations because they are almost
 		// always the same value for a given server.
-		if server1.pool.inUseConnectionCount() < server2.pool.inUseConnectionCount() {
+		if server1.OperationCount() < server2.OperationCount() {
 			return server1, nil
 		}
 		return server2, nil
@@ -678,7 +678,6 @@ func (t *Topology) processSRVResults(parsedHosts []string) bool {
 	t.subLock.Unlock()
 
 	return true
-
 }
 
 // apply updates the Topology and its underlying FSM based on the provided server description and returns the server

x/mongo/driver/topology/topology_test.go

@@ -722,7 +722,7 @@ func runInWindowTest(t *testing.T, directory string, filename string) {
 	// For each server state in the test's "mocked_topology_state", set the connection pool's
 	// in-use connections count to the test operation count value.
 	for _, state := range test.MockedTopologyState {
-		servers[state.Address].pool.inUseConnections = state.OperationCount
+		servers[state.Address].operationCount = state.OperationCount
 	}
 
 	// Create a new Topology, set the state to "connected", store a topology description
@@ -775,6 +775,8 @@ func runInWindowTest(t *testing.T, directory string, filename string) {
 		}
 
 		// Otherwise, check if the expected frequency is within the given tolerance range.
+		// TODO(GODRIVER-2179): Use assert.Deltaf() when we migrate all test code to the
+		// "testify/assert" or an API-compatible library for assertions.
 		low := expected - test.Outcome.Tolerance
 		high := expected + test.Outcome.Tolerance
 		assert.True(