lib/mongo/operation/shared/executable.rb in mongo-2.17.4 vs lib/mongo/operation/shared/executable.rb in mongo-2.18.0.beta1

@@ -25,31 +25,46 @@
     module Executable
 
       include ResponseHandling
 
       def do_execute(connection, context, options = {})
-        unpin_maybe(session) do
+        session&.materialize_if_needed
+        unpin_maybe(session, connection) do
           add_error_labels(connection, context) do
             add_server_diagnostics(connection) do
               get_result(connection, context, options).tap do |result|
                 if session
                   if session.in_transaction? &&
                     connection.description.load_balancer?
                   then
-                    if session.pinned_service_id
-                      unless session.pinned_service_id == connection.service_id
-                        raise Error::InternalDriverError, "Expected operation to use service #{session.pinned_session_id} but it used #{connection.service_id}"
+                    if session.pinned_connection_global_id
+                      unless session.pinned_connection_global_id == connection.global_id
+                        raise(
+                          Error::InternalDriverError,
+                          "Expected operation to use connection #{session.pinned_connection_global_id} but it used #{connection.global_id}"
+                        )
                       end
                     else
-                      session.pin_to_service(connection.service_id)
+                      session.pin_to_connection(connection.global_id)
+                      connection.pin
                     end
                   end
 
                   if session.snapshot? && !session.snapshot_timestamp
                     session.snapshot_timestamp = result.snapshot_timestamp
                   end
                 end
+
+                if result.has_cursor_id? &&
+                  connection.description.load_balancer?
+                then
+                  if result.cursor_id == 0
+                    connection.unpin
+                  else
+                    connection.pin
+                  end
+                end
                 process_result(result, connection)
               end
             end
           end
         end
@@ -80,10 +95,10 @@
       # Returns a Protocol::Message or nil as reply.
       def dispatch_message(connection, context, options = {})
         message = build_message(connection, context)
         message = message.maybe_encrypt(connection, context)
         reply = connection.dispatch([ message ], context, options)
-        [reply, connection.description]
+        [reply, connection.description, connection.global_id]
       end
 
       # @param [ Mongo::Server::Connection ] connection The connection on which
       #   the operation is performed.
       # @param [ Mongo::Operation::Context ] context The operation context.