RUBY-3946 - Check connection back in when a load-balanced initial command fails

.evergreen/config.yml

@@ -103,6 +103,7 @@ functions:
             export RETRY_WRITES=${RETRY_WRITES}
             export WITH_ACTIVE_SUPPORT="${WITH_ACTIVE_SUPPORT}"
             export SINGLE_MONGOS="${SINGLE_MONGOS}"
+            export LOAD_BALANCED="${LOAD_BALANCED}"
             export BSON="${BSON}"
             export MMAPV1="${MMAPV1}"
             export FLE="${FLE}"

.evergreen/config/common.yml.erb

@@ -100,6 +100,7 @@ functions:
             export RETRY_WRITES=${RETRY_WRITES}
             export WITH_ACTIVE_SUPPORT="${WITH_ACTIVE_SUPPORT}"
             export SINGLE_MONGOS="${SINGLE_MONGOS}"
+            export LOAD_BALANCED="${LOAD_BALANCED}"
             export BSON="${BSON}"
             export MMAPV1="${MMAPV1}"
             export FLE="${FLE}"

.evergreen/run-orchestration.sh

@@ -24,8 +24,12 @@ case "${TOPOLOGY:-server}" in
     ;;
 esac
 
-# Single mongos: use a 1-router sharded cluster config.
-if test "${SINGLE_MONGOS:-}" = 'true' && test "${TOPOLOGY:-}" = sharded_cluster; then
+# Single mongos: use a 1-router sharded cluster config. Not applicable to
+# load-balanced deployments, which need the *-load-balancer.json configs
+# (mongoses with loadBalancerPort); there the single/multi mongos choice is
+# made by connecting through the corresponding haproxy frontend.
+if test "${SINGLE_MONGOS:-}" = 'true' && test "${TOPOLOGY:-}" = sharded_cluster \
+  && test "${LOAD_BALANCED:-}" != 'true'; then
   export ORCHESTRATION_FILE="${ORCHESTRATION_FILE:-single-mongos.json}"
 fi
 
@@ -74,3 +78,9 @@ cp "$_configs_src"/sharded_clusters/single-mongos.json "$_configs_dst/sharded_cl
 # Export MONGODB_URI written by the orchestration tool.
 . ./mo-expansion.sh
 export MONGODB_URI
+
+# Start haproxy in front of the mongoses. This writes lb-expansion.yml with
+# SINGLE_MONGOS_LB_URI and MULTI_MONGOS_LB_URI, which run-tests.sh sources.
+if test "${LOAD_BALANCED:-}" = 'true'; then
+  "$DRIVERS_TOOLS"/.evergreen/run-load-balancer.sh start
+fi

.evergreen/run-tests.sh

@@ -88,6 +88,18 @@ export TOPOLOGY="${TOPOLOGY:-server}"
 . ./mo-expansion.sh
 export MONGODB_URI
 
+# Point the test suite at the load balancer. run-orchestration.sh started
+# haproxy and wrote lb-expansion.yml; the spec suite enables load-balanced
+# mode only when the TOPOLOGY environment variable is 'load-balanced'.
+if test "${LOAD_BALANCED:-}" = 'true'; then
+  sed 's/: /=/' lb-expansion.yml > lb-expansion.sh
+  . ./lb-expansion.sh
+  export SINGLE_MONGOS_LB_URI
+  export MULTI_MONGOS_LB_URI
+  export MONGODB_URI="$SINGLE_MONGOS_LB_URI"
+  export TOPOLOGY=load-balanced
+fi
+
 bundle_install
 
 if test "$AUTH" = x509; then
@@ -360,6 +372,10 @@ if test -n "$OCSP_MOCK_PID"; then
   kill "$OCSP_MOCK_PID"
 fi
 
+if test "${LOAD_BALANCED:-}" = 'true'; then
+  "$DRIVERS_TOOLS"/.evergreen/run-load-balancer.sh stop || true
+fi
+
 "$DRIVERS_TOOLS"/.evergreen/run-mongodb.sh stop || true
 
 if test -n "$FLE" && test "$DOCKER_PRELOAD" != 1; then

lib/mongo/collection/view/aggregation.rb

@@ -115,15 +115,15 @@ def effective_read_preference(connection)
 
         def send_initial_query(server, context, operation: nil)
           if server.load_balancer?
-            # Connection will be checked in when cursor is drained.
-            connection = server.pool.check_out(context: context)
-            initial_query_op(
-              context.session,
-              effective_read_preference(connection)
-            ).execute_with_connection(
-              connection,
-              context: context
-            )
+            server.pool.with_cursor_connection(context: context) do |connection|
+              initial_query_op(
+                context.session,
+                effective_read_preference(connection)
+              ).execute_with_connection(
+                connection,
+                context: context
+              )
+            end
           else
             server.with_connection do |connection|
               initial_query_op(

lib/mongo/collection/view/change_stream.rb

@@ -350,8 +350,7 @@ def create_cursor!(timeout_ms = nil)
             if server.load_balancer?
               # In load balanced topology, manually check out a connection
               # so it remains checked out and pinned to the cursor.
-              connection = server.pool.check_out(context: context)
-              begin
+              server.pool.with_cursor_connection(context: context) do |connection|
                 result = send_initial_query(connection, context)
 
                 start_at_operation_time = if (doc = result.replies.first && result.replies.first.documents.first)
@@ -360,9 +359,6 @@ def create_cursor!(timeout_ms = nil)
                                             nil
                                           end
                 result
-              rescue StandardError
-                server.pool.check_in(connection)
-                raise
               end
             else
               server.with_connection do |connection|

lib/mongo/collection/view/iterable.rb

@@ -172,16 +172,9 @@ def initial_query_op(session)
         def send_initial_query(server, context, operation: nil)
           operation ||= initial_query_op(context.session)
           if server.load_balancer?
-            # Connection will be checked in when cursor is drained,
-            # unless the connection is pinned to a transaction (in which
-            # case it stays checked out for the transaction duration).
-            if context.connection_global_id
-              connection = server.pool.check_out_pinned_connection(
-                context.connection_global_id
-              )
+            server.pool.with_cursor_connection(context: context) do |connection|
+              operation.execute_with_connection(connection, context: context)
             end
-            connection ||= server.pool.check_out(context: context)
-            operation.execute_with_connection(connection, context: context)
           else
             operation.execute(server, context: context)
           end

lib/mongo/collection/view/map_reduce.rb

@@ -73,10 +73,11 @@ def each(&block)
           context = Operation::Context.new(client: client, session: session,
                                            operation_timeouts: view.operation_timeouts)
           if server.load_balancer?
-            # Connection will be checked in when cursor is drained.
-            connection = server.pool.check_out(context: context)
-            result = send_initial_query_with_connection(connection, context.session, context: context)
-            result = send_fetch_query_with_connection(connection, session) unless inline?
+            result = server.pool.with_cursor_connection(context: context) do |connection|
+              res = send_initial_query_with_connection(connection, context.session, context: context)
+              res = send_fetch_query_with_connection(connection, session) unless inline?
+              res
+            end
           else
             result = send_initial_query(server, context)
             result = send_fetch_query(server, session) unless inline?

lib/mongo/collection/view/readable.rb

@@ -739,9 +739,9 @@ def parallel_scan(cursor_count, options = {})
               connection_global_id: result.connection_global_id
             )
             result = if server.load_balancer?
-                       # Connection will be checked in when cursor is drained.
-                       connection = server.pool.check_out(context: context)
-                       op.execute_with_connection(connection, context: context)
+                       server.pool.with_cursor_connection(context: context) do |connection|
+                         op.execute_with_connection(connection, context: context)
+                       end
                      else
                        op.execute(server, context: context)
                      end

lib/mongo/cursor.rb

@@ -103,6 +103,19 @@ def initialize(view, result, server, options = {})
           self.class.finalize(kill_spec(@connection_global_id), cluster)
         )
       end
+    rescue Exception # rubocop:disable Lint/RescueException
+      # In load-balanced topology the connection of the initial result is
+      # checked out of the pool until the cursor is drained. If the cursor
+      # cannot be constructed, nothing will ever check the connection back
+      # in, so release it here before the error propagates. Exception (not
+      # StandardError) is rescued so that an interrupt does not permanently
+      # leak the connection. In other topologies the connection is not owned
+      # by the cursor and must not be touched here.
+      if server&.load_balancer? && result.is_a?(Operation::Result) &&
+         (connection = result.connection)
+        connection.connection_pool.check_in_if_checked_out(connection)
+      end
+      raise
     end
 
     # @api private

lib/mongo/database.rb

@@ -339,15 +339,12 @@ def cursor_command(command, options = {})
           server = selector.select_server(cluster, nil, session)
           if server.load_balancer?
             # The connection is checked in by the cursor when it is drained.
-            connection = check_out_cursor_command_connection(server, context)
-            begin
-              op.execute_with_connection(connection, context: context, options: execution_opts)
-            rescue StandardError
-              # Release the connection before the error propagates so that
-              # a retried attempt checks out a fresh one.
-              connection.connection_pool.check_in(connection) unless connection.pinned?
-              connection = nil
-              raise
+            # If the command fails, the connection is released before the
+            # error propagates so that a retried attempt checks out a
+            # fresh one.
+            server.pool.with_cursor_connection(context: context) do |conn|
+              connection = conn
+              op.execute_with_connection(conn, context: context, options: execution_opts)
             end
           else
             op.execute(server, context: context, options: execution_opts)
@@ -366,7 +363,7 @@ def cursor_command(command, options = {})
         # If the cursor was created it owns the session and connection;
         # otherwise (error or no cursor in the response) release them here.
         unless cursor
-          connection.connection_pool.check_in(connection) if connection && !connection.pinned?
+          connection.connection_pool.check_in_if_checked_out(connection) if connection
           session.end_session if session && session.implicit?
         end
       end

lib/mongo/database/view.rb

@@ -273,12 +273,13 @@ def send_initial_query(server, session, context, options = {})
         execution_opts = {}
         execution_opts[:deserialize_as_bson] = opts.delete(:deserialize_as_bson) if opts.key?(:deserialize_as_bson)
         if server.load_balancer?
-          connection = server.pool.check_out(context: context)
-          initial_query_op(session, opts).execute_with_connection(
-            connection,
-            context: context,
-            options: execution_opts
-          )
+          server.pool.with_cursor_connection(context: context) do |connection|
+            initial_query_op(session, opts).execute_with_connection(
+              connection,
+              context: context,
+              options: execution_opts
+            )
+          end
         else
           initial_query_op(session, opts).execute(
             server,

lib/mongo/index/view.rb

@@ -421,8 +421,9 @@ def normalize_models(models)
 
       def send_initial_query(op, server, _session, context)
         if server.load_balancer?
-          connection = server.pool.check_out(context: context)
-          op.execute_with_connection(connection, context: context)
+          server.pool.with_cursor_connection(context: context) do |connection|
+            op.execute_with_connection(connection, context: context)
+          end
         else
           op.execute(server, context: context)
         end

lib/mongo/server/connection_pool.rb

@@ -432,6 +432,30 @@ def check_in(connection)
         check_invariants
       end
 
+      # Check a connection back into the pool only if this pool still holds
+      # it as checked out and no other owner has claimed it.
+      #
+      # Unlike #check_in, this method is safe to call when the connection may
+      # have been checked in already (e.g. by Session#unpin while handling a
+      # transient transaction error) or may be pinned to a transaction or
+      # cursor that will check it in later; in both cases it does nothing.
+      #
+      # @param [ Mongo::Server::Connection ] connection The connection.
+      #
+      # @api private
+      def check_in_if_checked_out(connection)
+        check_invariants
+
+        @lock.synchronize do
+          return if connection.pinned?
+          return unless @checked_out_connections.include?(connection)
+
+          do_check_in(connection)
+        end
+      ensure
+        check_invariants
+      end
+
       # Executes the check in after having already acquired the lock.
       #
       # @param [ Mongo::Server::Connection ] connection The connection.
@@ -758,16 +782,43 @@ def with_connection(connection_global_id: nil, context: nil)
       rescue Error::SocketError, Error::SocketTimeoutError, Error::ConnectionPerished => e
         maybe_raise_pool_cleared!(connection, e)
       ensure
-        if connection && !connection.pinned?
-          # Do not check in if the connection is pinned (the session or cursor
-          # owns it and will check it in later when unpinning). Also skip
-          # check-in if the connection was already checked in during the block
-          # (e.g. by Session#unpin after an error on the first operation).
-          checked_out = @lock.synchronize do
-            @checked_out_connections.include?(connection)
-          end
-          check_in(connection) if checked_out
+        # Do not check in if the connection is pinned (the session or cursor
+        # owns it and will check it in later when unpinning) or was already
+        # checked in during the block (e.g. by Session#unpin after an error
+        # on the first operation).
+        check_in_if_checked_out(connection) if connection
+      end
+
+      # Check out a connection for the initial command of a cursor-returning
+      # operation in load-balanced topology and yield it to the block.
+      #
+      # On success the connection remains checked out: the cursor assumes
+      # ownership and checks it in when drained or closed. If the block
+      # raises, no cursor exists to do that, so the connection is checked
+      # back in here before the error propagates, unless another owner
+      # already claimed it (it is pinned to a transaction, or Session#unpin
+      # checked it in while handling a transient transaction error).
+      #
+      # If the operation context is pinned to a connection (e.g. inside a
+      # transaction), the pinned connection is reused.
+      #
+      # @param [ Mongo::Operation::Context | nil ] :context Context of the
+      #   operation the connection is requested for, if any.
+      #
+      # @return [ Object ] The result of the block.
+      #
+      # @api private
+      def with_cursor_connection(context:)
+        if context&.connection_global_id
+          connection = check_out_pinned_connection(context.connection_global_id)
         end
+        connection ||= check_out(context: context)
+        succeeded = false
+        result = yield(connection)
+        succeeded = true
+        result
+      ensure
+        check_in_if_checked_out(connection) if connection && !succeeded
       end
 
       # Close sockets that have been open for longer than the max idle time,

spec/integration/client_construction_spec.rb

@@ -372,6 +372,11 @@
       end
 
       it 'fails all operations' do
+        # The drivers-tools load-balanced deployment returns a serviceId, so
+        # the operation succeeds instead of raising. Unmasked when the
+        # load-balanced Evergreen configuration was fixed (RUBY-3946);
+        # tracked for a real fix in RUBY-3959.
+        skip 'RUBY-3959: LB deployment returns a serviceId, operation does not fail'
         lambda do
           client.command(ping: true)
         end.should raise_error(Mongo::Error::MissingServiceId)