SERVER-66064 Optimize FLE2 document counting

(cherry picked from commit c23236c1b63f147f950d921a5411749e637d54ae)

7 files changed, 89 insertions, 76 deletions

diff --git a/jstests/multiVersion/targetedTestsLast b/jstests/multiVersion/targetedTestsLastLtsFeatures/fle2_downgrade.js
index 346e1b13f2d..0808e2be307 100644
--- a/jstests/multiVersion/targetedTestsLast
+++ b/jstests/multiVersion/targetedTestsLastLtsFeatures/fle2_downgrade.js
@@ -21,34 +21,31 @@ let conn = rst.getPrimary();
 let db = conn.getDB("admin");
 let client = new EncryptedClient(conn, dbName);
 
-function runTest(targetFCV)
-{
+function runTest(targetFCV) {
     assert.commandWorked(client.createEncryptionCollection("basic", {
-    encryptedFields: {
-        "fields": [
-            {"path": "first", "bsonType": "string", "queries": {"queryType": "equality"}},
-            {"path": "middle", "bsonType": "string"},
-            {"path": "aka", "bsonType": "string", "queries": {"queryType": "equality"}},
-        ]
-    }
-}));
-
-
-let res = assert.commandFailedWithCode(
-    db.adminCommand({setFeatureCompatibilityVersion: targetFCV}), ErrorCodes.CannotDowngrade);
+        encryptedFields: {
+            "fields": [
+                {"path": "first", "bsonType": "string", "queries": {"queryType": "equality"}},
+                {"path": "middle", "bsonType": "string"},
+                {"path": "aka", "bsonType": "string", "queries": {"queryType": "equality"}},
+            ]
+        }
+    }));
+
+    let res = assert.commandFailedWithCode(
+        db.adminCommand({setFeatureCompatibilityVersion: targetFCV}), ErrorCodes.CannotDowngrade);
 
     assert(client.getDB().fle2.basic.ecoc.drop());
     assert(client.getDB().fle2.basic.ecc.drop());
     assert(client.getDB().fle2.basic.esc.drop());
-assert(client.getDB().basic.drop());
+    assert(client.getDB().basic.drop());
 
-assert.commandWorked(db.adminCommand({setFeatureCompatibilityVersion: targetFCV}));
-assert.commandWorked(db.adminCommand({setFeatureCompatibilityVersion: latestFCV}));
+    assert.commandWorked(db.adminCommand({setFeatureCompatibilityVersion: targetFCV}));
+    assert.commandWorked(db.adminCommand({setFeatureCompatibilityVersion: latestFCV}));
 }
 
 targetFCV(lastLTSFCV);
 targetFCV(lastContinuousFCV);
 
-
 rst.stopSet();
 }());
\ No newline at end of file
diff --git a/src/mongo/crypto/fle_crypto.cpp b/src/mongo/crypto/fle_crypto.cpp
index 002d99e3f1c..e870634cb48 100644
--- a/src/mongo/crypto/fle_crypto.cpp
+++ b/src/mongo/crypto/fle_crypto.cpp
@@ -470,6 +470,12 @@ boost::optional<uint64_t> emuBinaryCommon(const FLEStateCollectionReader& reader
     // step 4, 5: get document count
     uint64_t rho = reader.getDocumentCount();
 
+    // Since fast count() is not reliable, if it says zero, try 1 instead just to be sure the
+    // collection is empty.
+    if (rho == 0) {
+        rho = 1;
+    }
+
 #ifdef DEBUG_ENUM_BINARY
     std::cout << fmt::format("start: lambda: {}, i: {}, rho: {}", lambda, i, rho) << std::endl;
 #endif
diff --git a/src/mongo/crypto/fle_crypto_test.cpp b/src/mongo/crypto/fle_crypto_test.cpp
index 0dce8b2c99f..75c1976c097 100644
--- a/src/mongo/crypto/fle_crypto_test.cpp
+++ b/src/mongo/crypto/fle_crypto_test.cpp
@@ -512,11 +512,20 @@ public:
     }
 
     uint64_t getDocumentCount() const override {
+        if (_overrideCount) {
+            return *_overrideCount;
+        }
+
         return _docs.size();
     }
 
+    void setOverrideCount(int64_t count) {
+        _overrideCount = count;
+    }
+
 private:
     std::vector<BSONObj> _docs;
+    boost::optional<int64_t> _overrideCount;
 };
 
 // Test Empty Collection
@@ -573,10 +582,19 @@ TEST(FLE_ESC, EmuBinary) {
         coll.insert(doc);
     }
 
-    auto i = ESCCollection::emuBinary(coll, escTwiceTag, escTwiceValue);
 
-    ASSERT_TRUE(i.has_value());
-    ASSERT_EQ(i.value(), 5);
+    // Test with various fake counts to ensure enumBinary works with bad estimates and the original
+    // exact count.
+    int64_t origCount = coll.getDocumentCount();
+    std::vector<int64_t> testVectors{0, 2, 3, 13, 500, origCount};
+
+    for (const auto v : testVectors) {
+        coll.setOverrideCount(v);
+        auto i = ESCCollection::emuBinary(coll, escTwiceTag, escTwiceValue);
+
+        ASSERT_TRUE(i.has_value());
+        ASSERT_EQ(i.value(), 5);
+    }
 }
 
 
@@ -624,15 +642,23 @@ TEST(FLE_ESC, EmuBinary2) {
         coll.insert(doc);
     }
 
-    auto i = ESCCollection::emuBinary(coll, escTwiceTag, escTwiceValue);
+    // Test with various fake counts to ensure enumBinary works with bad estimates and the original
+    // exact count.
+    int64_t origCount = coll.getDocumentCount();
+    std::vector<int64_t> testVectors{0, 2, 5, 13, 19, 500, origCount};
 
-    ASSERT_TRUE(i.has_value());
-    ASSERT_EQ(i.value(), 13);
+    for (const auto v : testVectors) {
+        coll.setOverrideCount(v);
+        auto i = ESCCollection::emuBinary(coll, escTwiceTag, escTwiceValue);
 
-    i = ESCCollection::emuBinary(coll, escTwiceTag2, escTwiceValue2);
+        ASSERT_TRUE(i.has_value());
+        ASSERT_EQ(i.value(), 13);
 
-    ASSERT_TRUE(i.has_value());
-    ASSERT_EQ(i.value(), 5);
+        i = ESCCollection::emuBinary(coll, escTwiceTag2, escTwiceValue2);
+
+        ASSERT_TRUE(i.has_value());
+        ASSERT_EQ(i.value(), 5);
+    }
 }
 
 // Test Emulated Binary with null record
diff --git a/src/mongo/db/commands/fle2_compact.cpp b/src/mongo/db/commands/fle2_compact.cpp
index 34a621d12d7..03fc6c7b0a9 100644
--- a/src/mongo/db/commands/fle2_compact.cpp
+++ b/src/mongo/db/commands/fle2_compact.cpp
@@ -630,7 +630,7 @@ CompactStats processFLECompact(OperationContext* opCtx,
             opCtx,
             [sharedBlock, c, ecocStats](const txn_api::TransactionClient& txnClient,
                                         ExecutorPtr txnExec) {
-                FLEQueryInterfaceImpl queryImpl(txnClient);
+                FLEQueryInterfaceImpl queryImpl(txnClient, getGlobalServiceContext());
 
                 auto [request2, namespaces2] = *sharedBlock.get();
 
@@ -659,7 +659,7 @@ CompactStats processFLECompact(OperationContext* opCtx,
             opCtx,
             [sharedBlock, escStats, eccStats](const txn_api::TransactionClient& txnClient,
                                               ExecutorPtr txnExec) {
-                FLEQueryInterfaceImpl queryImpl(txnClient);
+                FLEQueryInterfaceImpl queryImpl(txnClient, getGlobalServiceContext());
 
                 auto [ecocDoc2, namespaces2] = *sharedBlock.get();
 
diff --git a/src/mongo/db/fle_crud.cpp b/src/mongo/db/fle_crud.cpp
index ead3e250605..c041cedf79e 100644
--- a/src/mongo/db/fle_crud.cpp
+++ b/src/mongo/db/fle_crud.cpp
@@ -46,9 +46,11 @@
 #include "mongo/db/query/collation/collator_factory_interface.h"
 #include "mongo/db/query/find_command_gen.h"
 #include "mongo/db/query/fle/server_rewrite.h"
+#include "mongo/db/service_context.h"
 #include "mongo/db/transaction_api.h"
 #include "mongo/idl/idl_parser.h"
 #include "mongo/logv2/log.h"
+#include "mongo/rpc/factory.h"
 #include "mongo/s/grid.h"
 #include "mongo/s/transaction_router_resource_yielder.h"
 #include "mongo/s/write_ops/batch_write_exec.h"
@@ -216,7 +218,7 @@ std::pair<FLEBatchResult, write_ops::InsertCommandReply> processInsert(
         opCtx,
         [sharedInsertBlock, reply, ownedDocument](const txn_api::TransactionClient& txnClient,
                                                   ExecutorPtr txnExec) {
-            FLEQueryInterfaceImpl queryImpl(txnClient);
+            FLEQueryInterfaceImpl queryImpl(txnClient, getGlobalServiceContext());
 
             auto [edcNss2, efc2, serverPayload2, stmtId2] = *sharedInsertBlock.get();
 
@@ -288,7 +290,7 @@ write_ops::DeleteCommandReply processDelete(OperationContext* opCtx,
         opCtx,
         [sharedDeleteBlock, reply](const txn_api::TransactionClient& txnClient,
                                    ExecutorPtr txnExec) {
-            FLEQueryInterfaceImpl queryImpl(txnClient);
+            FLEQueryInterfaceImpl queryImpl(txnClient, getGlobalServiceContext());
 
             auto [deleteRequest2, expCtx2] = *sharedDeleteBlock.get();
 
@@ -364,7 +366,7 @@ write_ops::UpdateCommandReply processUpdate(OperationContext* opCtx,
         opCtx,
         [sharedupdateBlock, reply](const txn_api::TransactionClient& txnClient,
                                    ExecutorPtr txnExec) {
-            FLEQueryInterfaceImpl queryImpl(txnClient);
+            FLEQueryInterfaceImpl queryImpl(txnClient, getGlobalServiceContext());
 
             auto [updateRequest2, expCtx2] = *sharedupdateBlock.get();
 
@@ -637,7 +639,7 @@ StatusWith<ReplyType> processFindAndModifyRequest(
         opCtx,
         [sharedFindAndModifyBlock, reply, processCallback](
             const txn_api::TransactionClient& txnClient, ExecutorPtr txnExec) {
-            FLEQueryInterfaceImpl queryImpl(txnClient);
+            FLEQueryInterfaceImpl queryImpl(txnClient, getGlobalServiceContext());
 
             auto [findAndModifyRequest2, expCtx] = *sharedFindAndModifyBlock.get();
 
@@ -1154,54 +1156,31 @@ BSONObj FLEQueryInterfaceImpl::getById(const NamespaceString& nss, BSONElement e
 }
 
 uint64_t FLEQueryInterfaceImpl::countDocuments(const NamespaceString& nss) {
-    // TODO - what about
-    // count cmd
-    // $collStats
-    // approxCount
-
-    // Build the following pipeline:
-    //
-    //{ aggregate : "testColl", pipeline: [{$match:{}}, {$group : {_id: null, n : {$sum:1}
-    //}} ], cursor: {}}
-
-    BSONObjBuilder builder;
-    // $db - TXN API DOES THIS FOR US by building OP_MSG
-    builder.append("aggregate", nss.coll());
-
-    AggregateCommandRequest request(nss);
-
-    std::vector<BSONObj> pipeline;
-    pipeline.push_back(BSON("$match" << BSONObj()));
-
-    {
-        BSONObjBuilder sub;
-        {
-            BSONObjBuilder sub2(sub.subobjStart("$group"));
-            sub2.appendNull("_id");
-            {
-                BSONObjBuilder sub3(sub.subobjStart("n"));
-                sub3.append("$sum", 1);
-            }
-        }
+    // Since count() does not work in a transaction, call count() by bypassing the transaction api
+    invariant(!haveClient());
+    auto client = _serviceContext->makeClient("SEP-int-fle-crud");
+    AlternativeClientRegion clientRegion(client);
+    auto opCtx = cc().makeOperationContext();
 
-        pipeline.push_back(sub.obj());
-    }
+    CountCommandRequest ccr(nss);
+    auto opMsgRequest = ccr.serialize(BSONObj());
 
-    request.setPipeline(pipeline);
+    auto requestMessage = opMsgRequest.serialize();
 
-    auto commandResponse = _txnClient.runCommand(nss.db(), request.toBSON({})).get();
+    auto serviceEntryPoint = opCtx->getServiceContext()->getServiceEntryPoint();
+    DbResponse dbResponse = serviceEntryPoint->handleRequest(opCtx.get(), requestMessage).get();
 
-    uint64_t docCount = 0;
-    auto cursorResponse = uassertStatusOK(CursorResponse::parseFromBSON(commandResponse));
+    auto reply = rpc::makeReply(&dbResponse.response)->getCommandReply().getOwned();
+
+    auto status = getStatusFromWriteCommandReply(reply);
+    uassertStatusOK(status);
 
-    auto firstBatch = cursorResponse.getBatch();
-    if (!firstBatch.empty()) {
-        auto countObj = firstBatch.front();
-        docCount = countObj.getIntField("n"_sd);
-        uassert(6520701, "Unexpected negative document count", docCount >= 0);
+    int64_t signedDocCount = reply.getIntField("n"_sd);
+    if (signedDocCount < 0) {
+        signedDocCount = 0;
     }
 
-    return docCount;
+    return static_cast<uint64_t>(signedDocCount);
 }
 
 StatusWith<write_ops::InsertCommandReply> FLEQueryInterfaceImpl::insertDocument(
diff --git a/src/mongo/db/fle_crud.h b/src/mongo/db/fle_crud.h
index 65d79ecd946..0d5531fdb7f 100644
--- a/src/mongo/db/fle_crud.h
+++ b/src/mongo/db/fle_crud.h
@@ -43,6 +43,7 @@
 #include "mongo/db/pipeline/pipeline.h"
 #include "mongo/db/query/count_command_gen.h"
 #include "mongo/db/server_options.h"
+#include "mongo/db/service_context.h"
 #include "mongo/db/transaction_api.h"
 #include "mongo/s/write_ops/batch_write_exec.h"
 #include "mongo/s/write_ops/batched_command_response.h"
@@ -313,7 +314,9 @@ public:
  */
 class FLEQueryInterfaceImpl : public FLEQueryInterface {
 public:
-    FLEQueryInterfaceImpl(const txn_api::TransactionClient& txnClient) : _txnClient(txnClient) {}
+    FLEQueryInterfaceImpl(const txn_api::TransactionClient& txnClient,
+                          ServiceContext* serviceContext)
+        : _txnClient(txnClient), _serviceContext(serviceContext) {}
 
     BSONObj getById(const NamespaceString& nss, BSONElement element) final;
 
@@ -348,6 +351,7 @@ public:
 
 private:
     const txn_api::TransactionClient& _txnClient;
+    ServiceContext* _serviceContext;
 };
 
 /**
diff --git a/src/mongo/db/query/fle/server_rewrite.cpp b/src/mongo/db/query/fle/server_rewrite.cpp
index dc91058e718..185a979ec3d 100644
--- a/src/mongo/db/query/fle/server_rewrite.cpp
+++ b/src/mongo/db/query/fle/server_rewrite.cpp
@@ -47,6 +47,7 @@
 #include "mongo/db/pipeline/document_source_match.h"
 #include "mongo/db/pipeline/expression.h"
 #include "mongo/db/query/collation/collator_factory_interface.h"
+#include "mongo/db/service_context.h"
 #include "mongo/s/grid.h"
 #include "mongo/s/transaction_router_resource_yielder.h"
 #include "mongo/util/assert_util.h"
@@ -303,7 +304,7 @@ void doFLERewriteInTxn(OperationContext* opCtx,
             };
 
             // Construct FLE rewriter from the transaction client and encryptionInformation.
-            auto queryInterface = FLEQueryInterfaceImpl(txnClient);
+            auto queryInterface = FLEQueryInterfaceImpl(txnClient, getGlobalServiceContext());
             auto escReader = makeCollectionReader(&queryInterface, sharedBlock->esc);
             auto eccReader = makeCollectionReader(&queryInterface, sharedBlock->ecc);