lib/mongo/operation/shared/executable.rb in mongo-2.15.1 vs lib/mongo/operation/shared/executable.rb in mongo-2.16.0.alpha1

@@ -29,10 +29,23 @@
       def do_execute(connection, context, options = {})
         unpin_maybe(session) 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}"
+                      end
+                    else
+                      session.pin_to_service(connection.service_id)
+                    end
+                  end
+                end
                 process_result(result, connection)
               end
             end
           end
         end
@@ -91,11 +104,11 @@
         result
       end
 
       def process_result_for_sdam(result, connection)
         if (result.not_master? || result.node_recovering?) &&
-          connection.generation >= connection.server.pool.generation
+          connection.generation >= connection.server.pool.generation(service_id: connection.service_id)
         then
           if result.node_shutting_down?
             keep_pool = false
           else
             # Max wire version needs to be examined while the server is known
@@ -103,9 +116,10 @@
           end
 
           connection.server.unknown!(
             keep_connection_pool: keep_pool,
             generation: connection.generation,
+            service_id: connection.service_id,
             topology_version: result.topology_version,
           )
 
           connection.server.scan_semaphore.signal
         end