SERVER-66533 update FLE array tag size estimation

3 files changed, 109 insertions, 6 deletions

diff --git a/src/mongo/crypto/fle_tags.cpp b/src/mongo/crypto/fle_tags.cpp
index 7ec5ec0c83b..a0de37b2f42 100644
--- a/src/mongo/crypto/fle_tags.cpp
+++ b/src/mongo/crypto/fle_tags.cpp
@@ -38,6 +38,48 @@ namespace mongo::fle {
 using DerivedToken = FLEDerivedFromDataTokenAndContentionFactorTokenGenerator;
 using TwiceDerived = FLETwiceDerivedTokenGenerator;
 
+size_t sizeArrayElementsMemory(size_t tagCount);
+
+namespace {
+
+inline constexpr size_t arrayElementSize(int digits) {
+    constexpr size_t sizeOfType = 1;
+    constexpr size_t sizeOfBinDataLength = 4;
+    constexpr size_t sizeOfNullMarker = 1;
+    constexpr size_t sizeOfSubType = 1;
+    constexpr size_t sizeOfData = sizeof(PrfBlock);
+    return sizeOfType + sizeOfBinDataLength + sizeOfNullMarker + digits + sizeOfSubType +
+        sizeOfData;
+}
+
+void verifyTagsWillFit(size_t tagCount, size_t memoryLimit) {
+    constexpr size_t largestElementSize = arrayElementSize(std::numeric_limits<size_t>::digits10);
+    constexpr size_t ridiculousNumberOfTags =
+        std::numeric_limits<size_t>::max() / largestElementSize;
+    uassert(6653300, "Encrypted rewrite too many tags", tagCount < ridiculousNumberOfTags);
+    uassert(6401800,
+            "Encrypted rewrite memory limit exceeded",
+            sizeArrayElementsMemory(tagCount) <= memoryLimit);
+}
+
+}  // namespace
+
+size_t sizeArrayElementsMemory(size_t tagCount) {
+    size_t size = 0;
+    size_t power = 1;
+    size_t digits = 1;
+    size_t accountedTags = 0;
+    while (tagCount >= power) {
+        power *= 10;
+        size_t count = std::min(tagCount, power) - accountedTags;
+        size += arrayElementSize(digits) * count;
+        accountedTags += count;
+        digits++;
+    }
+    return size;
+}
+
+
 // The algorithm for constructing a list of tags matching an equality predicate on an encrypted
 // field is as follows:
 //
@@ -135,9 +177,8 @@ std::vector<PrfBlock> readTagsWithContention(const FLEStateCollectionReader& esc
         return acc + eccDoc.end - eccDoc.start + 1;
     });
     auto cumTagCount = binaryTags.size() + numInserts - numDeletes;
-    uassert(6401800,
-            "Encrypted rewrite memory limit exceeded",
-            cumTagCount * sizeof(PrfBlock) <= memoryLimit);
+
+    verifyTagsWillFit(cumTagCount, memoryLimit);
 
     for (uint64_t i = 1; i <= numInserts; i++) {
         if (auto it = std::lower_bound(
diff --git a/src/mongo/db/fle_crud_test.cpp b/src/mongo/db/fle_crud_test.cpp
index c4835e57772..527dd5bca11 100644
--- a/src/mongo/db/fle_crud_test.cpp
+++ b/src/mongo/db/fle_crud_test.cpp
@@ -70,6 +70,11 @@
 #include "mongo/util/uuid.h"
 
 namespace mongo {
+
+namespace fle {
+size_t sizeArrayElementsMemory(size_t tagCount);
+}
+
 namespace {
 
 constexpr auto kIndexKeyId = "12345678-1234-9876-1234-123456789012"_sd;
@@ -1180,8 +1185,8 @@ TEST_F(FleTagsTest, MemoryLimit) {
 
     const auto tagLimit = 10;
 
-    // Set memory limit to 10 tags * 32 bytes per tag
-    internalQueryFLERewriteMemoryLimit.store(tagLimit * 32);
+    // Set memory limit to 10 tags * 40 bytes per tag
+    internalQueryFLERewriteMemoryLimit.store(tagLimit * 40);
 
     // Do 10 inserts
     for (auto i = 0; i < tagLimit; i++) {
@@ -1201,5 +1206,60 @@ TEST_F(FleTagsTest, MemoryLimit) {
     // readTags returns 10 tags which does not exceed memory limit.
     ASSERT_EQ(tagLimit, readTags(doc).size());
 }
+
+TEST_F(FleTagsTest, SampleMemoryLimit) {
+
+    struct S {
+        size_t count;
+        size_t size;
+    };
+
+    // clang-format off
+    static const std::vector<S> testVector{
+        { 0, 0 },
+        { 1, 40 },
+        { 5, 200 },
+        { 10, 400 },
+        { 11, 441 },
+        { 98, 4008 },
+        { 99, 4049 },
+        { 100, 4090 },
+        { 101, 4132 },
+        { 219, 9088 },
+        { 944, 39538 },
+        { 998, 41806 },
+        { 999, 41848 },
+        { 1000, 41890 },
+        { 1001, 41933 },
+        { 1025, 42965 },
+        { 1498, 63304 },
+        { 2049, 86997 },
+        { 2907, 123891 },
+        { 5232, 223866 },
+        { 5845, 250225 },
+        { 7203, 308619 },
+        { 7786, 333688 },
+        { 8383, 359359 },
+        { 9171, 393243 },
+        { 9974, 427772 },
+        { 9986, 428288 },
+        { 9998, 428804 },
+        { 9999, 428847 },
+        { 10000, 428890 },
+        { 10001, 428934 },
+        { 10056, 431354 },
+        { 10907, 468798 },
+        { 12500, 538890 },
+        { 13778, 595122 },
+        { 13822, 597058 },
+    };
+    // clang-format on
+
+    for (auto& xp : testVector) {
+        auto size = mongo::fle::sizeArrayElementsMemory(xp.count);
+        ASSERT_EQ(xp.size, size);
+    }
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_knobs.idl b/src/mongo/db/query/query_knobs.idl
index 7e2fee1563f..2e8f4d7358c 100644
--- a/src/mongo/db/query/query_knobs.idl
+++ b/src/mongo/db/query/query_knobs.idl
@@ -852,7 +852,8 @@ server_parameters:
     on_update: plan_cache_util::clearSbeCacheOnParameterChange
 
   internalQueryFLERewriteMemoryLimit:
-    description: "Maximum memory available for encrypted field query rewrites."
+    description: "Maximum memory available for encrypted field query rewrites in bytes. Must be
+    more than zero and less than 16Mb"
     set_at: [ startup, runtime ]
     cpp_varname: "internalQueryFLERewriteMemoryLimit"
     cpp_vartype: AtomicWord<int>
@@ -860,6 +861,7 @@ server_parameters:
         expr: 14 * 1024 * 1024
     validator:
         gt: 0
+        lt: 16777216
 
 # Note for adding additional query knobs:
 #