SERVER-124376: GeoNear should not spill if allow disk use is set to false (#52044)

GitOrigin-RevId: 7e713e0881b3acdd4dfddce5bf1ee0321c40330a

etc/backports_required_for_multiversion_tests.yml

@@ -33,6 +33,8 @@ last-continuous:
       ticket: SERVER-121533
     - test_file: jstests/core/timeseries/query/timeseries_lookup_index_use.js
       ticket: SERVER-95352
+    - test_file: jstests/core/index/geo/geo_circle_spilling.js
+      ticket: SERVER-124376
   suites: null
 last-lts:
   all:
@@ -706,4 +708,6 @@ last-lts:
       ticket: SERVER-95352
     - test_file: jstests/sharding/resharding_change_stream_internal_ops.js
       ticket: SERVER-120962
+    - test_file: jstests/core/index/geo/geo_circle_spilling.js
+      ticket: SERVER-124376
   suites: null

jstests/core/index/geo/geo_circle_spilling.js

@@ -33,6 +33,12 @@ const nearPredicate = {
 
 function assertSpillingAndAllDocumentsReturned(coll) {
     jsTest.log.info("Running query", nearPredicate);
+
+    assert.throwsWithCode(
+        () => coll.find(nearPredicate).allowDiskUse(false).itcount(),
+        ErrorCodes.QueryExceededMemoryLimitNoDiskUseAllowed,
+    );
+
     const explain = coll.find(nearPredicate).explain("executionStats");
     let assertedSpilling = false;
     for (let stages of getExecutionStages(explain)) {

src/mongo/db/exec/classic/near.cpp

@@ -58,7 +58,7 @@ NearStage::NearStage(ExpressionContext* expCtx,
           makeSortOptions(),
           SorterKeyComparator{},
           NoOpBound{},
-          (feature_flags::gFeatureFlagExtendedAutoSpilling.isEnabled())
+          (expCtx->getAllowDiskUse() && feature_flags::gFeatureFlagExtendedAutoSpilling.isEnabled())
               ? std::make_shared<
                     sorter::FileBasedSpiller<SorterKey, SorterValue, SorterKeyComparator>>(
                     expCtx->getTempDir(),
@@ -212,7 +212,7 @@ PlanStage::StageState NearStage::bufferNext(WorkingSetID* toReturn) {
     _memoryTracker.set(_seenDocuments.getApproximateSize() + _resultBuffer.stats().memUsage());
     if (!_memoryTracker.withinMemoryLimit() &&
         feature_flags::gFeatureFlagExtendedAutoSpilling.isEnabled()) {
-        spill(loadMemoryLimit(StageMemoryLimit::NearStageMaxMemoryBytes));
+        spill(_memoryTracker.maxAllowedMemoryUsageBytes());
     }
     _specificStats.peakTrackedMemBytes = _memoryTracker.peakTrackedMemoryBytes();
     uassert(12227900, "Near stage exceeded memory limit", _memoryTracker.withinMemoryLimit());
@@ -261,29 +261,24 @@ PlanStage::StageState NearStage::advanceNext(WorkingSetID* toReturn) {
 }
 
 SortOptions NearStage::makeSortOptions() {
-    if (feature_flags::gFeatureFlagExtendedAutoSpilling.isEnabled()) {
-        return SortOptions{}  // Spilling will handled externally by NearStage::spill method
-            .MaxMemoryUsageBytes(std::numeric_limits<int64_t>::max());
-    } else {
-        return SortOptions{}.MaxMemoryUsageBytes(std::numeric_limits<int64_t>::max());
-    }
+    return SortOptions{}  // Spilling will handled externally by NearStage::spill method
+        .MaxMemoryUsageBytes(std::numeric_limits<int64_t>::max());
 }
 
 void NearStage::updateSpillingStats() {
-
     auto additionalSpilledBytes = _sorterFileStats.bytesSpilledUncompressed() -
         _specificStats.spillingStats.getSpilledBytes();
+    auto additionalSpilledRecords = _resultBuffer.stats().spilledKeyValuePairs() -
+        _specificStats.spillingStats.getSpilledRecords();
 
-    auto spilledDataStorageIncrease = _specificStats.spillingStats.updateSpillingStats(
-        1 /*spills*/,
-        additionalSpilledBytes,
-        _resultBuffer.stats().spilledKeyValuePairs(),
-        _sorterFileStats.bytesSpilled());
+    auto spilledDataStorageIncrease =
+        _specificStats.spillingStats.updateSpillingStats(1 /*spills*/,
+                                                         additionalSpilledBytes,
+                                                         additionalSpilledRecords,
+                                                         _sorterFileStats.bytesSpilled());
 
-    geoNearCounters.incrementPerSpilling(1,
-                                         additionalSpilledBytes,
-                                         _resultBuffer.stats().spilledKeyValuePairs(),
-                                         spilledDataStorageIncrease);
+    geoNearCounters.incrementPerSpilling(
+        1, additionalSpilledBytes, additionalSpilledRecords, spilledDataStorageIncrease);
 }
 
 void NearStage::spill(uint64_t maxMemoryBytes) {
@@ -292,6 +287,11 @@ void NearStage::spill(uint64_t maxMemoryBytes) {
     if (totalMemoryUsage <= maxMemoryBytes) {
         return;
     }
+    uassert(ErrorCodes::QueryExceededMemoryLimitNoDiskUseAllowed,
+            str::stream() << _commonStats.stageTypeStr
+                          << " stage exceeded memory limit and can't spill to disk. Set "
+                             "allowDiskUse: true to allow spilling",
+            expCtx()->getAllowDiskUse());
     _resultBuffer.forceSpill();
     _memoryTracker.set(_seenDocuments.getApproximateSize() + _resultBuffer.stats().memUsage());
     updateSpillingStats();

src/mongo/dbtests/query_stage_near.cpp

@@ -303,6 +303,12 @@ TEST_F(QueryStageNearTest, Spilling) {
 
     const auto* stats = static_cast<const NearStats*>(nearStage.getSpecificStats());
     ASSERT_GT(stats->spillingStats.getSpills(), 0);
+    ASSERT_GT(stats->spillingStats.getSpilledRecords(), results.size() / 2);
+
+    size_t approximateBytesPerSpilledRecord =
+        stats->spillingStats.getSpilledBytes() / stats->spillingStats.getSpilledRecords();
+    ASSERT_LTE(16, approximateBytesPerSpilledRecord);
+    ASSERT_GTE(64, approximateBytesPerSpilledRecord);
 }
 
 TEST_F(QueryStageNearTest, MemoryTracking) {