RUBY-2520 Change estimatedDocumentCount() to use the $collStats Agg Stage Instead of Count Command (#2198)

Co-authored-by: Oleg Pudeyev <oleg@bsdpower.com>

Author: p-mongo

.evergreen/config.yml

@@ -928,7 +928,7 @@ buildvariants:
     tasks:
       - name: "test-mlaunch"
 
-  - matrix_name: "mongo-4.9"
+  - matrix_name: "mongo-4.9-api-version"
     matrix_spec:
       auth-and-ssl: ["auth-and-ssl", "noauth-and-nossl"]
       ruby: "ruby-2.7"

.evergreen/config/standard.yml.erb

@@ -131,7 +131,7 @@ buildvariants:
     tasks:
       - name: "test-mlaunch"
 
-  - matrix_name: "mongo-4.9"
+  - matrix_name: "mongo-4.9-api-version"
     matrix_spec:
       auth-and-ssl: ["auth-and-ssl", "noauth-and-nossl"]
       ruby: "ruby-2.7"

lib/mongo/collection/view/readable.rb

@@ -230,24 +230,46 @@ def estimated_document_count(opts = {})
             raise ArgumentError, "Cannot call estimated_document_count when querying with a filter"
           end
 
-          cmd = { count: collection.name }
-          cmd[:maxTimeMS] = opts[:max_time_ms] if opts[:max_time_ms]
-          if read_concern
-            cmd[:readConcern] = Options::Mapper.transform_values_to_strings(
-              read_concern)
-          end
           Mongo::Lint.validate_underscore_read_preference(opts[:read])
           read_pref = opts[:read] || read_preference
           selector = ServerSelector.get(read_pref || server_selector)
           with_session(opts) do |session|
+            context = Operation::Context.new(client: client, session: session)
             read_with_retry(session, selector) do |server|
-              Operation::Count.new(
-                selector: cmd,
-                db_name: database.name,
-                read: read_pref,
-                session: session,
-              ).execute(server, context: Operation::Context.new(client: client, session: session))
-            end.n.to_i
+              if server.description.server_version_gte?('5.0')
+                pipeline = [
+                  {'$collStats' => {'count' => {}}},
+                  {'$group' => {'_id' => 1, 'n' => {'$sum' => '$count'}}},
+                ]
+                spec = Builder::Aggregation.new(pipeline, self, options.merge(session: session)).specification
+                result = Operation::Aggregate.new(spec).execute(server, context: context)
+                result.documents.first.fetch('n')
+              else
+                cmd = { count: collection.name }
+                cmd[:maxTimeMS] = opts[:max_time_ms] if opts[:max_time_ms]
+                if read_concern
+                  cmd[:readConcern] = Options::Mapper.transform_values_to_strings(
+                    read_concern)
+                end
+                result = Operation::Count.new(
+                  selector: cmd,
+                  db_name: database.name,
+                  read: read_pref,
+                  session: session,
+                ).execute(server, context: context)
+                result.n.to_i
+              end
+            end
+          end
+        rescue Error::OperationFailure => exc
+          if exc.code == 26
+            # NamespaceNotFound
+            # This should only happen with the aggregation pipeline path
+            # (server 4.9+). Previous servers should return 0 for nonexistent
+            # collections.
+            0
+          else
+            raise
           end
         end
 

lib/mongo/operation.rb

@@ -6,6 +6,7 @@
 require 'mongo/operation/shared/executable_no_validate'
 require 'mongo/operation/shared/executable_transaction_label'
 require 'mongo/operation/shared/polymorphic_lookup'
+require 'mongo/operation/shared/polymorphic_operation'
 require 'mongo/operation/shared/polymorphic_result'
 require 'mongo/operation/shared/read_preference_supported'
 require 'mongo/operation/shared/bypass_document_validation'
@@ -19,8 +20,6 @@
 require 'mongo/operation/shared/object_id_generator'
 require 'mongo/operation/shared/op_msg_or_command'
 require 'mongo/operation/shared/op_msg_or_find_command'
-require 'mongo/operation/shared/op_msg_or_list_indexes_command'
-require 'mongo/operation/shared/collections_info_or_list_collections'
 
 require 'mongo/operation/op_msg_base'
 require 'mongo/operation/command'

lib/mongo/operation/collections_info.rb

@@ -25,7 +25,24 @@ module Operation
     # @since 2.0.0
     class CollectionsInfo
       include Specifiable
-      include CollectionsInfoOrListCollections
+      include PolymorphicOperation
+      include PolymorphicLookup
+
+      private
+
+      def final_operation(connection)
+         op_class = if connection.features.list_collections_enabled?
+          if connection.features.op_msg_enabled?
+            ListCollections::OpMsg
+          else
+            ListCollections::Command
+          end
+        else
+          CollectionsInfo::Command
+        end
+
+        op_class.new(spec)
+      end
     end
   end
 end

lib/mongo/operation/indexes.rb

@@ -27,7 +27,21 @@ module Operation
     # @since 2.0.0
     class Indexes
       include Specifiable
-      include OpMsgOrListIndexesCommand
+      include PolymorphicOperation
+      include PolymorphicLookup
+
+      private
+
+      def final_operation(connection)
+        cls = if connection.features.op_msg_enabled?
+          polymorphic_class(self.class.name, :OpMsg)
+        elsif connection.features.list_indexes_enabled?
+          polymorphic_class(self.class.name, :Command)
+        else
+          polymorphic_class(self.class.name, :Legacy)
+        end
+        cls.new(spec)
+      end
     end
   end
 end

lib/mongo/operation/shared/op_msg_or_command.rb

@@ -20,15 +20,9 @@ module Operation
     #
     # @api private
     module OpMsgOrCommand
+      include PolymorphicOperation
       include PolymorphicLookup
 
-      def execute(server, context:, options: {})
-        server.with_connection do |connection|
-          operation = final_operation(connection)
-          operation.execute(connection, context: context, options: options)
-        end
-      end
-
       private
 
       def final_operation(connection)

lib/mongo/operation/shared/op_msg_or_find_command.rb

@@ -21,15 +21,9 @@ module Operation
     #
     # @api private
     module OpMsgOrFindCommand
+      include PolymorphicOperation
       include PolymorphicLookup
 
-      def execute(server, context:)
-        server.with_connection do |connection|
-          operation = final_operation(connection)
-          operation.execute(connection, context: context)
-        end
-      end
-
       private
 
       def final_operation(connection)

lib/mongo/operation/shared/op_msg_or_list_indexes_command.rb

@@ -1,4 +1,4 @@
-# Copyright (C) 2020 MongoDB Inc.
+# Copyright (C) 2021 MongoDB Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,39 +15,25 @@
 module Mongo
   module Operation
 
+    # Shared behavior of implementing an operation differently based on
+    # the server that will be executing the operation.
+    #
     # @api private
-    module CollectionsInfoOrListCollections
-      include PolymorphicLookup
+    module PolymorphicOperation
 
       # Execute the operation.
       #
       # @param [ Mongo::Server ] server The server to send the operation to.
       # @param [ Operation::Context ] context The operation context.
+      # @param [ Hash ] options Operation execution options.
       #
-      # @return [ Mongo::Operation::CollectionsInfo::Result,
-      #           Mongo::Operation::ListCollections::Result ] The operation result.
-      def execute(server, context:)
+      # @return [ Mongo::Operation::Result ] The operation result.
+      def execute(server, context:, options: {})
         server.with_connection do |connection|
           operation = final_operation(connection)
-          operation.execute(connection, context: context)
+          operation.execute(connection, context: context, options: options)
         end
       end
-
-      private
-
-      def final_operation(connection)
-         op_class = if connection.features.list_collections_enabled?
-          if connection.features.op_msg_enabled?
-            ListCollections::OpMsg
-          else
-            ListCollections::Command
-          end
-        else
-          CollectionsInfo::Command
-        end
-
-        op_class.new(spec)
-      end
     end
   end
 end