SERVER-69768 Include key pattern in range deletion task documents

9 files changed, 117 insertions, 15 deletions

diff --git a/src/mongo/db/s/migration_coordinator.cpp b/src/mongo/db/s/migration_coordinator.cpp
index 5844d30c3a5..feaf7252cc2 100644
--- a/src/mongo/db/s/migration_coordinator.cpp
+++ b/src/mongo/db/s/migration_coordinator.cpp
@@ -75,6 +75,7 @@ MigrationCoordinator::MigrationCoordinator(MigrationSessionId sessionId,
                                            UUID collectionUuid,
                                            ChunkRange range,
                                            ChunkVersion preMigrationChunkVersion,
+                                           const KeyPattern& shardKeyPattern,
                                            bool waitForDelete)
     : _migrationInfo(UUID::gen(),
                      std::move(sessionId),
@@ -86,6 +87,7 @@ MigrationCoordinator::MigrationCoordinator(MigrationSessionId sessionId,
                      std::move(recipientShard),
                      std::move(range),
                      std::move(preMigrationChunkVersion)),
+      _shardKeyPattern(shardKeyPattern),
       _waitForDelete(waitForDelete) {}
 
 MigrationCoordinator::MigrationCoordinator(const MigrationCoordinatorDocument& doc)
@@ -122,6 +124,7 @@ void MigrationCoordinator::startMigration(OperationContext* opCtx) {
                                         _waitForDelete ? CleanWhenEnum::kNow
                                                        : CleanWhenEnum::kDelayed);
     donorDeletionTask.setPending(true);
+    donorDeletionTask.setKeyPattern(*_shardKeyPattern);
     const auto currentTime = VectorClock::get(opCtx)->getTime();
     donorDeletionTask.setTimestamp(currentTime.clusterTime().asTimestamp());
     migrationutil::persistRangeDeletionTaskLocally(
diff --git a/src/mongo/db/s/migration_coordinator.h b/src/mongo/db/s/migration_coordinator.h
index 23355c4c96c..5c24400ef7c 100644
--- a/src/mongo/db/s/migration_coordinator.h
+++ b/src/mongo/db/s/migration_coordinator.h
@@ -49,6 +49,7 @@ public:
                          UUID collectionUuid,
                          ChunkRange range,
                          ChunkVersion preMigrationChunkVersion,
+                         const KeyPattern& shardKeyPattern,
                          bool waitForDelete);
 
     MigrationCoordinator(const MigrationCoordinatorDocument& doc);
@@ -123,6 +124,7 @@ private:
     void _waitForReleaseRecipientCriticalSectionFutureIgnoreShardNotFound(OperationContext* opCtx);
 
     MigrationCoordinatorDocument _migrationInfo;
+    boost::optional<KeyPattern> _shardKeyPattern;
     bool _waitForDelete = false;
     boost::optional<ExecutorFuture<void>> _releaseRecipientCriticalSectionFuture;
 };
diff --git a/src/mongo/db/s/migration_destination_manager.cpp b/src/mongo/db/s/migration_destination_manager.cpp
index 72ee678b386..435cb76e8f5 100644
--- a/src/mongo/db/s/migration_destination_manager.cpp
+++ b/src/mongo/db/s/migration_destination_manager.cpp
@@ -1283,6 +1283,7 @@ void MigrationDestinationManager::_migrateDriver(OperationContext* outerOpCtx,
             recipientDeletionTask.setPending(true);
             const auto currentTime = VectorClock::get(outerOpCtx)->getTime();
             recipientDeletionTask.setTimestamp(currentTime.clusterTime().asTimestamp());
+            recipientDeletionTask.setKeyPattern(KeyPattern(_shardKeyPattern));
 
             // It is illegal to wait for write concern with a session checked out, so persist the
             // range deletion task with an immediately satsifiable write concern and then wait for
diff --git a/src/mongo/db/s/migration_source_manager.cpp b/src/mongo/db/s/migration_source_manager.cpp
index 13ddda47064..45be08187fa 100644
--- a/src/mongo/db/s/migration_source_manager.cpp
+++ b/src/mongo/db/s/migration_source_manager.cpp
@@ -289,6 +289,7 @@ void MigrationSourceManager::startClone() {
                              *_collectionUUID,
                              ChunkRange(*_args.getMin(), *_args.getMax()),
                              *_chunkVersion,
+                             KeyPattern(metadata.getKeyPattern()),
                              _args.getWaitForDelete());
 
         _state = kCloning;
diff --git a/src/mongo/db/s/range_deleter_service.cpp b/src/mongo/db/s/range_deleter_service.cpp
index ee5ff61f608..e0027e09bfe 100644
--- a/src/mongo/db/s/range_deleter_service.cpp
+++ b/src/mongo/db/s/range_deleter_service.cpp
@@ -320,7 +320,8 @@ SharedSemiFuture<void> RangeDeleterService::registerTask(
             .then([this,
                    dbName = rdt.getNss().dbName(),
                    collectionUuid = rdt.getCollectionUuid(),
-                   range = rdt.getRange()]() {
+                   range = rdt.getRange(),
+                   optKeyPattern = rdt.getKeyPattern()]() {
                 return withTemporaryOperationContext(
                     [&](OperationContext* opCtx) {
                         // A task is considered completed when all the following conditions are met:
@@ -344,10 +345,16 @@ SharedSemiFuture<void> RangeDeleterService::registerTask(
                                             "range"_attr = redact(range.toString()));
 
                                         auto shardKeyPattern =
-                                            getShardKeyPattern(opCtx, dbName, collectionUuid);
-
-                                        uassertStatusOK(deleteRangeInBatches(
-                                            opCtx, dbName, collectionUuid, shardKeyPattern, range));
+                                            (optKeyPattern ? *optKeyPattern
+                                                           : getShardKeyPattern(
+                                                                 opCtx, dbName, collectionUuid));
+
+                                        uassertStatusOK(
+                                            deleteRangeInBatches(opCtx,
+                                                                 dbName,
+                                                                 collectionUuid,
+                                                                 shardKeyPattern.toBSON(),
+                                                                 range));
                                         orphansRemovalCompleted = true;
                                     } catch (ExceptionFor<ErrorCodes::NamespaceNotFound>&) {
                                         // No orphaned documents to remove from a dropped collection
diff --git a/src/mongo/db/s/range_deleter_service_test.cpp b/src/mongo/db/s/range_deleter_service_test.cpp
index 8d465cfc34d..51492211a12 100644
--- a/src/mongo/db/s/range_deleter_service_test.cpp
+++ b/src/mongo/db/s/range_deleter_service_test.cpp
@@ -61,13 +61,13 @@ void RangeDeleterServiceTest::setUp() {
         AutoGetCollection autoColl(opCtx, nsCollA, MODE_IX);
         uuidCollA = autoColl.getCollection()->uuid();
         nssWithUuid[uuidCollA] = nsCollA;
-        _setFilteringMetadataWithUUID(opCtx, uuidCollA);
+        _setFilteringMetadataByUUID(opCtx, uuidCollA);
     }
     {
         AutoGetCollection autoColl(opCtx, nsCollB, MODE_IX);
         uuidCollB = autoColl.getCollection()->uuid();
         nssWithUuid[uuidCollB] = nsCollB;
-        _setFilteringMetadataWithUUID(opCtx, uuidCollB);
+        _setFilteringMetadataByUUID(opCtx, uuidCollB);
     }
 
     rangeDeletionTask0ForCollA = createRangeDeletionTaskWithOngoingQueries(
@@ -85,8 +85,8 @@ void RangeDeleterServiceTest::tearDown() {
     ShardServerTestFixture::tearDown();
 }
 
-void RangeDeleterServiceTest::_setFilteringMetadataWithUUID(OperationContext* opCtx,
-                                                            const UUID& uuid) {
+void RangeDeleterServiceTest::_setFilteringMetadataByUUID(OperationContext* opCtx,
+                                                          const UUID& uuid) {
     const OID epoch = OID::gen();
     NamespaceString nss = nssWithUuid[uuid];
 
@@ -747,4 +747,69 @@ TEST_F(RangeDeleterServiceTest, OnlyRemoveDocumentsInRangeToDelete) {
     ASSERT_EQUALS(dbclient.count(NamespaceString::kRangeDeletionNamespace), 0);
 }
 
+TEST_F(RangeDeleterServiceTest, RegisterAndProcessSingleTaskWithKeyPattern) {
+    auto rds = RangeDeleterService::get(opCtx);
+    auto taskWithOngoingQueries =
+        createRangeDeletionTaskWithOngoingQueries(uuidCollA,
+                                                  BSON(kShardKey << 0),
+                                                  BSON(kShardKey << 10),
+                                                  CleanWhenEnum::kNow,
+                                                  false,
+                                                  KeyPattern(kShardKeyPattern));
+
+    auto completionFuture =
+        registerAndCreatePersistentTask(opCtx,
+                                        taskWithOngoingQueries->getTask(),
+                                        taskWithOngoingQueries->getOngoingQueriesFuture());
+
+    // The task can't be processed (hence completed) before ongoing queries drain
+    ASSERT(!completionFuture.isReady());
+    ASSERT_EQ(1, rds->getNumRangeDeletionTasksForCollection(uuidCollA));
+
+    // Make sure deletion can proceed even without filtering metadata
+    _clearFilteringMetadataByUUID(opCtx, uuidCollA);
+
+    // Pretend ongoing queries have drained and check task is processed
+    taskWithOngoingQueries->drainOngoingQueries();
+    completionFuture.get(opCtx);
+    ASSERT_EQ(0, rds->getNumRangeDeletionTasksForCollection(uuidCollA));
+}  // namespace mongo
+
+TEST_F(RangeDeleterServiceTest, PerformActualRangeDeletionWithKeyPattern) {
+    auto rds = RangeDeleterService::get(opCtx);
+    auto taskWithOngoingQueries =
+        createRangeDeletionTaskWithOngoingQueries(uuidCollA,
+                                                  BSON(kShardKey << 0),
+                                                  BSON(kShardKey << 10),
+                                                  CleanWhenEnum::kNow,
+                                                  false,
+                                                  KeyPattern(kShardKeyPattern));
+    auto nss = nssWithUuid[uuidCollA];
+    DBDirectClient dbclient(opCtx);
+
+    insertDocsWithinRange(opCtx, nss, 0, 10, 10);
+    ASSERT_EQUALS(dbclient.count(nss, BSONObj()), 10);
+
+    auto completionFuture =
+        registerAndCreatePersistentTask(opCtx,
+                                        taskWithOngoingQueries->getTask(),
+                                        taskWithOngoingQueries->getOngoingQueriesFuture());
+
+    // The task can't be processed (hence completed) before ongoing queries drain
+    ASSERT(!completionFuture.isReady());
+    ASSERT_EQ(1, rds->getNumRangeDeletionTasksForCollection(uuidCollA));
+    ASSERT_EQUALS(dbclient.count(NamespaceString::kRangeDeletionNamespace), 1);
+    verifyRangeDeletionTasks(opCtx, uuidCollA, {taskWithOngoingQueries->getTask().getRange()});
+
+    // Make sure deletion can proceed even without filtering metadata
+    _clearFilteringMetadataByUUID(opCtx, uuidCollA);
+
+    // Pretend ongoing queries have drained and check task is processed
+    taskWithOngoingQueries->drainOngoingQueries();
+    completionFuture.get(opCtx);
+    ASSERT_EQ(0, rds->getNumRangeDeletionTasksForCollection(uuidCollA));
+    ASSERT_EQUALS(dbclient.count(nss), 0);
+
+}  // namespace mongo
+
 }  // namespace mongo
diff --git a/src/mongo/db/s/range_deleter_service_test.h b/src/mongo/db/s/range_deleter_service_test.h
index 7c689178ab1..cc44ac03225 100644
--- a/src/mongo/db/s/range_deleter_service_test.h
+++ b/src/mongo/db/s/range_deleter_service_test.h
@@ -75,7 +75,7 @@ public:
     inline static const BSONObj kShardKeyPattern = BSON(kShardKey << 1);
 
 private:
-    void _setFilteringMetadataWithUUID(OperationContext* opCtx, const UUID& uuid);
+    void _setFilteringMetadataByUUID(OperationContext* opCtx, const UUID& uuid);
 
     // Scoped objects
     RAIIServerParameterControllerForTest enableFeatureFlag{"featureFlagRangeDeleterService", true};
@@ -87,14 +87,16 @@ RangeDeletionTask createRangeDeletionTask(const UUID& collectionUUID,
                                           const BSONObj& min,
                                           const BSONObj& max,
                                           CleanWhenEnum whenToClean = CleanWhenEnum::kNow,
-                                          bool pending = true);
+                                          bool pending = true,
+                                          boost::optional<KeyPattern> keyPattern = boost::none);
 
 std::shared_ptr<RangeDeletionWithOngoingQueries> createRangeDeletionTaskWithOngoingQueries(
     const UUID& collectionUUID,
     const BSONObj& min,
     const BSONObj& max,
     CleanWhenEnum whenToClean = CleanWhenEnum::kNow,
-    bool pending = true);
+    bool pending = true,
+    boost::optional<KeyPattern> keyPattern = boost::none);
 
 SharedSemiFuture<void> registerAndCreatePersistentTask(
     OperationContext* opCtx,
@@ -108,6 +110,9 @@ void verifyRangeDeletionTasks(OperationContext* opCtx,
                               UUID uuidColl,
                               std::vector<ChunkRange> expectedChunkRanges);
 
+/* Unset any filtering metadata associated with the specified collection */
+void _clearFilteringMetadataByUUID(OperationContext* opCtx, const UUID& uuid);
+
 // CRUD operation over `config.rangeDeletions`
 void insertRangeDeletionTaskDocument(OperationContext* opCtx, const RangeDeletionTask& rdt);
 void updatePendingField(OperationContext* opCtx, UUID rdtId, bool pending);
diff --git a/src/mongo/db/s/range_deleter_service_test_util.cpp b/src/mongo/db/s/range_deleter_service_test_util.cpp
index 6bfc2a76d5e..607f51abbf8 100644
--- a/src/mongo/db/s/range_deleter_service_test_util.cpp
+++ b/src/mongo/db/s/range_deleter_service_test_util.cpp
@@ -27,7 +27,9 @@
  *    it in the license file.
  */
 
+#include "mongo/db/catalog_raii.h"
 #include "mongo/db/persistent_task_store.h"
+#include "mongo/db/s/collection_sharding_runtime.h"
 #include "mongo/db/s/range_deleter_service_test.h"
 
 namespace mongo {
@@ -57,7 +59,8 @@ RangeDeletionTask createRangeDeletionTask(const UUID& collectionUUID,
                                           const BSONObj& min,
                                           const BSONObj& max,
                                           CleanWhenEnum whenToClean,
-                                          bool pending) {
+                                          bool pending,
+                                          boost::optional<KeyPattern> keyPattern) {
     RangeDeletionTask rdt;
     rdt.setId(UUID::gen());
     rdt.setNss(RangeDeleterServiceTest::nssWithUuid[collectionUUID]);
@@ -66,6 +69,8 @@ RangeDeletionTask createRangeDeletionTask(const UUID& collectionUUID,
     rdt.setRange(ChunkRange(min, max));
     rdt.setWhenToClean(whenToClean);
     rdt.setPending(pending);
+    rdt.setKeyPattern(keyPattern);
+
     return rdt;
 }
 
@@ -74,9 +79,10 @@ std::shared_ptr<RangeDeletionWithOngoingQueries> createRangeDeletionTaskWithOngo
     const BSONObj& min,
     const BSONObj& max,
     CleanWhenEnum whenToClean,
-    bool pending) {
+    bool pending,
+    boost::optional<KeyPattern> keyPattern) {
     return std::make_shared<RangeDeletionWithOngoingQueries>(
-        createRangeDeletionTask(collectionUUID, min, max, whenToClean, pending));
+        createRangeDeletionTask(collectionUUID, min, max, whenToClean, pending, keyPattern));
 }
 
 // TODO review this method: the task may be registered, finish and be recreated by inserting the
@@ -174,4 +180,11 @@ void verifyRangeDeletionTasks(OperationContext* opCtx,
     }
 }
 
+void _clearFilteringMetadataByUUID(OperationContext* opCtx, const UUID& uuid) {
+    NamespaceString nss = RangeDeleterServiceTest::nssWithUuid[uuid];
+
+    AutoGetCollection autoColl(opCtx, nss, LockMode::MODE_X);
+    CollectionShardingRuntime::get(opCtx, nss)->clearFilteringMetadata(opCtx);
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/s/range_deletion_task.idl b/src/mongo/db/s/range_deletion_task.idl
index f33b2b24463..45932f0c6a7 100644
--- a/src/mongo/db/s/range_deletion_task.idl
+++ b/src/mongo/db/s/range_deletion_task.idl
@@ -34,6 +34,7 @@ global:
 
 imports:
     - "mongo/db/basic_types.idl"
+    - "mongo/db/keypattern.idl"
     - "mongo/s/sharding_types.idl"
     - "mongo/db/session/logical_session_id.idl"
     - "mongo/s/chunk_range.idl"
@@ -85,3 +86,7 @@ structs:
                 type: safeInt64
                 description: "The estimated number of orphaned documents in the range"
                 default: 0
+            keyPattern:
+                type: KeyPattern
+                description: "The shard key pattern at the moment the range deletion task was created"
+                optional: true