RUBY-2682 RUBY-2443 test that driver closes cursors left open by server when result set is fully iterated (#2378)

* calculate remaining number of documents in result set for OP_MSG servers

* debugging

* RUBY-2682 cleaned up/added number_returned tests

* RUBY-2682 add test for nonzero cursor id

* RUBY-2682 fix msg test

* RUBY-2682 change number_returned default

* RUBY-2682 dont raise error in number_returned

* RUBY-2682 fix spec test for consistency (RUBY-2443)

* RUBY-2682 fix number_returned without batch test

* prohibit unexpected uses of returned_count

* adjust unit test

Co-authored-by: Oleg Pudeyev <[email protected]>

Author: neilshweky

lib/mongo/cursor.rb

@@ -383,6 +383,14 @@ def fully_iterated?
 
     private
 
+    def batch_size_for_get_more
+      if batch_size && use_limit?
+        [batch_size, @remaining].min
+      else
+        batch_size
+      end
+    end
+
     def exhausted?
       limited? ? @remaining <= 0 : false
     end
@@ -405,7 +413,7 @@ def get_more_operation
         cursor_id: id,
         # 3.2+ servers use batch_size, 3.0- servers use to_return.
         # TODO should to_return be calculated in the operation layer?
-        batch_size: batch_size,
+        batch_size: batch_size_for_get_more,
         to_return: to_return,
         max_time_ms: if view.respond_to?(:max_await_time_ms) &&
           view.max_await_time_ms &&

lib/mongo/operation/map_reduce/result.rb

@@ -131,6 +131,22 @@ def cursor_id
           0
         end
 
+        # Get the number of documents returned by the server in this batch.
+        #
+        # Map/Reduce operation returns documents inline without using
+        # cursors; as such, the standard Mongo::Reply#returned_count does
+        # not work correctly for Map/Reduce.
+        #
+        # Note that the Map/Reduce operation is limited to max BSON document
+        # size (16 MB) in its inline result set.
+        #
+        # @return [ Integer ] The number of documents returned.
+        #
+        # @api public
+        def returned_count
+          reply.documents.length
+        end
+
         private
 
         def first_document

lib/mongo/operation/result.rb

@@ -238,10 +238,7 @@ def reply
         end
       end
 
-      # Get the count of documents returned by the server.
-      #
-      # @example Get the number returned.
-      #   result.returned_count
+      # Get the number of documents returned by the server in this batch.
       #
       # @return [ Integer ] The number of documents returned.
       #

lib/mongo/protocol/msg.rb

@@ -298,6 +298,24 @@ def maybe_add_server_api(server_api)
         Msg.new(@flags, @options, main_document, *@sequences)
       end
 
+      # Returns the number of documents returned from the server.
+      #
+      # The Msg instance must be for a server reply and the reply must return
+      # an active cursor (either a newly created one or one whose iteration is
+      # continuing via getMore).
+      #
+      # @return [ Integer ] Number of returned documents.
+      def number_returned
+        if doc = documents.first
+          if cursor = doc['cursor']
+            if batch = cursor['firstBatch'] || cursor['nextBatch']
+              return batch.length
+            end
+          end
+        end
+        raise NotImplementedError, "number_returned is only defined for cursor replies"
+      end
+
       private
 
       # Validate that the documents in this message are all smaller than the

spec/mongo/cursor_spec.rb

@@ -543,6 +543,36 @@
         end
       end
     end
+
+    context 'when the result set is iterated fully and the cursor id is non-zero' do
+      min_server_fcv '5.0'
+
+      let(:documents) do
+        (1..5).map{ |i| { field: "test#{i}" }}
+      end
+
+      let(:view) { collection.find(test:{'$gte'=>BSON::MinKey.new}).sort(test:1).limit(5).batch_size(4) }
+
+      before do
+        view.to_a
+      end
+
+      it 'schedules a get more command' do
+        get_more_commands = subscriber.started_events.select { |e| e.command_name == 'getMore' }
+        expect(get_more_commands.length).to be 1
+      end
+
+      it 'has a non-zero cursor id on successful get more' do
+        get_more_commands = subscriber.succeeded_events.select { |e| e.command_name == 'getMore' }
+        expect(get_more_commands.length).to be 1
+        expect(get_more_commands[0].reply['cursor']['id']).to_not be 0
+      end
+
+      it 'schedules a kill cursors command' do
+        get_more_commands = subscriber.started_events.select { |e| e.command_name == 'killCursors' }
+        expect(get_more_commands.length).to be 1
+      end
+    end
   end
 
   describe '#inspect' do

spec/mongo/protocol/msg_spec.rb

@@ -525,4 +525,45 @@
       end
     end
   end
+
+  describe '#number_returned' do
+
+    let(:batch) do
+      (1..2).map{ |i| { field: "test#{i}" }}
+    end
+
+    context 'when the msg contains a find document' do
+
+      let(:find_document) { { "cursor" => { "firstBatch" => batch } } }
+
+      let(:find_message) do
+        described_class.new(flags, options, find_document, *sequences)
+      end
+
+      it 'returns the correct number_returned' do
+        expect(find_message.number_returned).to eq(2)
+      end
+    end
+
+    context 'when the msg contains a getmore document' do
+      let(:next_document) { { "cursor" => { "nextBatch" => batch } } }
+
+      let(:next_message) do
+        described_class.new(flags, options, next_document, *sequences)
+      end
+
+      it 'returns the correct number_returned' do
+        expect(next_message.number_returned).to eq(2)
+      end
+    end
+
+    context 'when the msg contains a document without first/nextBatch' do
+
+      it 'raises NotImplementedError' do
+        lambda do
+          message.number_returned
+        end.should raise_error(NotImplementedError, /number_returned is only defined for cursor replies/)
+      end
+    end
+  end
 end

spec/spec_tests/data/unified/valid-pass/poc-command-monitoring.yml

@@ -69,8 +69,7 @@ tests:
               command:
                 getMore: { $$type: [ int, long ] }
                 collection: *collection0Name
-                # https://jira.mongodb.org/browse/RUBY-2443
-                batchSize: 3
+                batchSize: 1
               commandName: getMore
               databaseName: *database0Name
           - commandSucceededEvent: