SERVER-68042 Registration of duplicate range deletion task must return original completion future

3 files changed, 35 insertions, 20 deletions

diff --git a/src/mongo/db/s/range_deleter_service.cpp b/src/mongo/db/s/range_deleter_service.cpp
index 9273e6e3e6c..df0b142791a 100644
--- a/src/mongo/db/s/range_deleter_service.cpp
+++ b/src/mongo/db/s/range_deleter_service.cpp
@@ -28,6 +28,9 @@
  */
 
 #include "mongo/db/s/range_deleter_service.h"
+#include "mongo/logv2/log.h"
+
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kShardingRangeDeleter
 
 namespace mongo {
 SharedSemiFuture<void> RangeDeleterService::registerTask(
@@ -40,7 +43,7 @@ SharedSemiFuture<void> RangeDeleterService::registerTask(
     initialFutures.push_back(blockUntilRegistered.getFuture().semi().thenRunOn(_executor));
     initialFutures.push_back(std::move(waitForActiveQueriesToComplete).thenRunOn(_executor));
 
-    auto completionFuture =
+    auto chainCompletionFuture =
         // Step 1: wait for the task to be registered on the service and for the draining of
         // ongoing queries that are retaining the orphaned range
         whenAllSucceed(std::move(initialFutures))
@@ -74,29 +77,34 @@ SharedSemiFuture<void> RangeDeleterService::registerTask(
             .semi()
             .share();
 
-    bool inserted = [&]() {
+    auto [taskCompletionFuture, inserted] = [&]() -> std::pair<SharedSemiFuture<void>, bool> {
         stdx::lock_guard<Latch> lg(_mutex);
-        auto [_, inserted] = _rangeDeletionTasks[rdt.getCollectionUuid()].insert(
-            RangeDeletion(rdt, completionFuture));
-        return inserted;
+        auto [registeredTask, inserted] = _rangeDeletionTasks[rdt.getCollectionUuid()].insert(
+            std::make_shared<RangeDeletion>(RangeDeletion(rdt, chainCompletionFuture)));
+        auto retFuture = static_cast<RangeDeletion*>(registeredTask->get())->getCompletionFuture();
+        return {retFuture, inserted};
     }();
 
-    if (!inserted) {
+    if (inserted) {
+        // The range deletion task has been registered, so the chain execution can be unblocked
+        blockUntilRegistered.setFrom(Status::OK());
+    } else {
         // Tried to register a duplicate range deletion task: invalidate the chain
         auto errStatus =
             Status(ErrorCodes::Error(67635), "Not scheduling duplicated range deletion");
+        LOGV2_WARNING(6804200,
+                      "Tried to register duplicate range deletion task. This results in a no-op.",
+                      "collectionUUID"_attr = rdt.getCollectionUuid(),
+                      "range"_attr = rdt.getRange());
         blockUntilRegistered.setFrom(errStatus);
-        return blockUntilRegistered.getFuture();
     }
 
-    // The range deletion task has been registered, so the chain execution can be unblocked
-    blockUntilRegistered.setFrom(Status::OK());
-    return completionFuture;
+    return taskCompletionFuture;
 }
 
 void RangeDeleterService::deregisterTask(const UUID& collUUID, const ChunkRange& range) {
     stdx::lock_guard<Latch> lg(_mutex);
-    _rangeDeletionTasks[collUUID].erase(range);
+    _rangeDeletionTasks[collUUID].erase(std::make_shared<ChunkRange>(range));
 }
 
 int RangeDeleterService::getNumRangeDeletionTasksForCollection(const UUID& collectionUUID) {
diff --git a/src/mongo/db/s/range_deleter_service.h b/src/mongo/db/s/range_deleter_service.h
index 3d090dfda6a..19336603d47 100644
--- a/src/mongo/db/s/range_deleter_service.h
+++ b/src/mongo/db/s/range_deleter_service.h
@@ -69,7 +69,7 @@ private:
             : ChunkRange(task.getRange().getMin(), task.getRange().getMax()),
               _completion(completion) {}
 
-        SharedSemiFuture<void> getCompletionFuture() {
+        SharedSemiFuture<void> getCompletionFuture() const {
             return _completion;
         }
 
@@ -86,13 +86,14 @@ private:
      * sharded collection).
      */
     struct RANGES_COMPARATOR {
-        bool operator()(const ChunkRange& a, const ChunkRange& b) const {
-            return a.getMin().woCompare(b.getMin()) < 0;
+        bool operator()(const std::shared_ptr<ChunkRange>& a,
+                        const std::shared_ptr<ChunkRange>& b) const {
+            return a->getMin().woCompare(b->getMin()) < 0;
         }
     };
 
     // Keeping track of per-collection registered range deletion tasks
-    stdx::unordered_map<UUID, std::set<ChunkRange, RANGES_COMPARATOR>, UUID::Hash>
+    stdx::unordered_map<UUID, std::set<std::shared_ptr<ChunkRange>, RANGES_COMPARATOR>, UUID::Hash>
         _rangeDeletionTasks;
 
     // Mono-threaded executor processing range deletion tasks
diff --git a/src/mongo/db/s/range_deleter_service_test.cpp b/src/mongo/db/s/range_deleter_service_test.cpp
index 14e9328f37d..caea1a20e54 100644
--- a/src/mongo/db/s/range_deleter_service_test.cpp
+++ b/src/mongo/db/s/range_deleter_service_test.cpp
@@ -109,22 +109,28 @@ TEST_F(RangeDeleterServiceTest, RegisterAndProcessSingleTask) {
     ASSERT_EQ(0, rds.getNumRangeDeletionTasksForCollection(uuidCollA));
 }
 
-TEST_F(RangeDeleterServiceTest, CantRegisterDuplicateTaskForSameRange) {
+TEST_F(RangeDeleterServiceTest, RegisterDuplicateTaskForSameRangeReturnsOriginalFuture) {
     RangeDeleterService rds;
     auto* taskWithOngoingQueries = &rangeDeletionTask0ForCollA;
 
     auto originalTaskCompletionFuture = rds.registerTask(
         taskWithOngoingQueries->getTask(), taskWithOngoingQueries->getOngoingQueriesFuture());
 
-    // Trying registering a duplicate task must return a future with error
+    // Trying registering a duplicate task must return a future without throwing errors
     auto duplicateTaskCompletionFuture =
         rds.registerTask(taskWithOngoingQueries->getTask(), SemiFuture<void>::makeReady());
-    ASSERT(duplicateTaskCompletionFuture.isReady());
-    ASSERT_THROWS_CODE(duplicateTaskCompletionFuture.get(opCtx), DBException, 67635);
 
-    // Trying registering a duplicate task did not affect the state of the original task
+    // Check that the state of the original task is not affected
     ASSERT(!originalTaskCompletionFuture.isReady());
+
+    // Check that no new task has been registered
     ASSERT_EQ(1, rds.getNumRangeDeletionTasksForCollection(uuidCollA));
+
+    // Check that - upon range deletion completion - both original and "duplicate" futures are ready
+    taskWithOngoingQueries->drainOngoingQueries();
+    originalTaskCompletionFuture.get(opCtx);
+    ASSERT(duplicateTaskCompletionFuture.isReady());
+    ASSERT_EQ(0, rds.getNumRangeDeletionTasksForCollection(uuidCollA));
 }
 
 TEST_F(RangeDeleterServiceTest, RegisterAndProcessMoreTasksForSameCollection) {