Fix RUBY-2675 QueryCache returns wrong results after partial iteration of result set, e.g. after calling none? (#2329)

Co-authored-by: Oleg Pudeyev <code@olegp.name>

Author: p-mongo

docs/reference/query-cache.txt

@@ -68,6 +68,42 @@ And to disable the query cache in the context of a block:
 You may check whether the query cache is enabled at any time by calling
 ``Mongo::QueryCache.enabled?``, which will return ``true`` or ``false``.
 
+
+Interactions With Fibers
+========================
+
+The Query cache enablement flag is stored in fiber-local storage (using
+`Thread.current <https://ruby-doc.org/core/Thread.html#class-Thread-label-Fiber-local+vs.+Thread-local>`_.
+This, in principle, permits query cache state to be per fiber, although
+this is not currently tested.
+
+There are methods in the Ruby standard library, like ``Enumerable#next``,
+that `utilize fibers <https://stackoverflow.com/questions/11057223/how-does-rubys-enumerator-object-iterate-externally-over-an-internal-iterator/11058270#11058270>`_
+in their implementation. These methods would not see the query cache
+enablement flag when it is set by the applications, and subsequently would
+not use the query cache. For example, the following code does not utilize
+the query cache despite requesting it:
+
+.. code-block:: ruby
+
+    Mongo::QueryCache.enabled = true
+    
+    client['artists'].find({}, limit: 1).to_enum.next
+    # Issues the query again.
+    client['artists'].find({}, limit: 1).to_enum.next
+
+Rewriting this code to use ``first`` instead of ``next`` would make it use
+the query cache:
+
+.. code-block:: ruby
+
+    Mongo::QueryCache.enabled = true
+    
+    client['artists'].find({}, limit: 1).first
+    # Utilizes the cached result from the first query.
+    client['artists'].find({}, limit: 1).first
+
+
 .. _query-cache-matching:
 
 Query Matching

lib/mongo/collection/view/iterable.rb

@@ -38,18 +38,28 @@ module Iterable
         #
         # @yieldparam [ Hash ] Each matching document.
         def each
-          @cursor = if use_query_cache? && cached_cursor
+          # If the caching cursor is closed and was not fully iterated,
+          # the documents we have in it are not the complete result set and
+          # we have no way of completing that iteration.
+          # Therefore, discard that cursor and start iteration again.
+          # The case of the caching cursor not being closed and not having
+          # been fully iterated isn't tested - see RUBY-2773.
+          @cursor = if use_query_cache? && cached_cursor && (
+            cached_cursor.fully_iterated? || !cached_cursor.closed?
+          )
             cached_cursor
           else
             session = client.send(:get_session, @options)
-            select_cursor(session)
+            select_cursor(session).tap do |cursor|
+              if use_query_cache?
+                # No need to store the cursor in the query cache if there is
+                # already a cached cursor stored at this key.
+                QueryCache.set(cursor, **cache_options)
+              end
+            end
           end
 
           if use_query_cache?
-            # No need to store the cursor in the query cache if there is
-            # already a cached cursor stored at this key.
-            QueryCache.set(@cursor, **cache_options) unless cached_cursor
-
             # If a query with a limit is performed, the query cache will
             # re-use results from an earlier query with the same or larger
             # limit, and then impose the lower limit during iteration.

lib/mongo/cursor.rb

@@ -212,10 +212,12 @@ def try_next
         unless closed?
           if exhausted?
             close
+            @fully_iterated = true
             raise StopIteration
           end
           @documents = get_more
         else
+          @fully_iterated = true
           raise StopIteration
         end
       else
@@ -232,6 +234,9 @@ def try_next
       # over the last document, or if the batch is empty
       if @documents.size <= 1
         cache_batch_resume_token
+        if closed?
+          @fully_iterated = true
+        end
       end
 
       return @documents.shift
@@ -352,6 +357,11 @@ def get_more
       process(get_more_operation.execute(@server, context: Operation::Context.new(client: client, session: @session)))
     end
 
+    # @api private
+    def fully_iterated?
+      !!@fully_iterated
+    end
+
     private
 
     def exhausted?

spec/integration/query_cache_spec.rb

@@ -1045,4 +1045,49 @@
       expect(events.length).to eq(2)
     end
   end
+
+  context 'when result set has multiple documents and cursor is iterated partially' do
+
+    before do
+      Mongo::QueryCache.enabled = false
+      5.times do
+        authorized_collection.insert_one({ name: 'testing' })
+      end
+    end
+
+    shared_examples 'retrieves full result set on second iteration' do
+      it 'retrieves full result set on second iteration' do
+        Mongo::QueryCache.clear
+        Mongo::QueryCache.enabled = true
+
+        partial_first_iteration
+
+        authorized_collection.find.to_a.length.should == 5
+      end
+
+    end
+
+    context 'using each & break' do
+      let(:partial_first_iteration) do
+        called = false
+        authorized_collection.find.each do
+          called = true
+          break
+        end
+        called.should be true
+      end
+
+      include_examples 'retrieves full result set on second iteration'
+    end
+
+    context 'using next' do
+      let(:partial_first_iteration) do
+        # #next is executed in its own fiber, and query cache is disabled
+        # for that operation.
+        authorized_collection.find.to_enum.next
+      end
+
+      include_examples 'retrieves full result set on second iteration'
+    end
+  end
 end