SERVER-60465 Align BalancerCommandsScheduler API with defragmentation algorithm

5 files changed, 79 insertions, 59 deletions

diff --git a/dump_python3.22527.core b/dump_python3.22527.core
new file mode 100644
index 00000000000..31b4c9b561b
--- /dev/null
+++ b/dump_python3.22527.core
diff --git a/src/mongo/db/s/balancer/balancer_commands_scheduler.h b/src/mongo/db/s/balancer/balancer_commands_scheduler.h
index c749f970f68..2083972413b 100644
--- a/src/mongo/db/s/balancer/balancer_commands_scheduler.h
+++ b/src/mongo/db/s/balancer/balancer_commands_scheduler.h
@@ -37,7 +37,7 @@
 namespace mongo {
 
 class ChunkType;
-class ShardKeyPattern;
+class KeyPattern;
 class MigrationSecondaryThrottleOptions;
 
 /**
@@ -150,29 +150,31 @@ public:
     virtual std::unique_ptr<MergeChunksResponse> requestMergeChunks(
         OperationContext* opCtx,
         const NamespaceString& nss,
-        const ChunkType& lowerBound,
-        const ChunkType& upperBound) = 0;
+        const ShardId& shardId,
+        const ChunkRange& chunkRange,
+        const ChunkVersion& version) = 0;
 
     virtual std::unique_ptr<SplitVectorResponse> requestSplitVector(
         OperationContext* opCtx,
         const NamespaceString& nss,
         const ChunkType& chunk,
-        const ShardKeyPattern& shardKeyPattern,
+        const KeyPattern& keyPattern,
         const SplitVectorSettings& commandSettings) = 0;
 
     virtual std::unique_ptr<SplitChunkResponse> requestSplitChunk(
         OperationContext* opCtx,
         const NamespaceString& nss,
         const ChunkType& chunk,
-        const ShardKeyPattern& shardKeyPattern,
+        const KeyPattern& keyPattern,
         const std::vector<BSONObj>& splitPoints) = 0;
 
-    virtual std::unique_ptr<ChunkDataSizeResponse> requestChunkDataSize(
-        OperationContext* opCtx,
-        const NamespaceString& nss,
-        const ChunkType& chunk,
-        const ShardKeyPattern& shardKeyPattern,
-        bool estimatedValue) = 0;
+    virtual std::unique_ptr<ChunkDataSizeResponse> requestDataSize(OperationContext* opCtx,
+                                                                   const NamespaceString& nss,
+                                                                   const ShardId& shardId,
+                                                                   const ChunkRange& chunkRange,
+                                                                   const ChunkVersion& version,
+                                                                   const KeyPattern& keyPattern,
+                                                                   bool estimatedValue) = 0;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/s/balancer/balancer_commands_scheduler_impl.cpp b/src/mongo/db/s/balancer/balancer_commands_scheduler_impl.cpp
index 602423605b5..efc2ec9278b 100644
--- a/src/mongo/db/s/balancer/balancer_commands_scheduler_impl.cpp
+++ b/src/mongo/db/s/balancer/balancer_commands_scheduler_impl.cpp
@@ -140,18 +140,12 @@ std::unique_ptr<MoveChunkResponse> BalancerCommandsSchedulerImpl::requestMoveChu
 std::unique_ptr<MergeChunksResponse> BalancerCommandsSchedulerImpl::requestMergeChunks(
     OperationContext* opCtx,
     const NamespaceString& nss,
-    const ChunkType& lowerBound,
-    const ChunkType& upperBound) {
-    invariant(lowerBound.getShard() == upperBound.getShard());
-    invariant(lowerBound.getMax().woCompare(upperBound.getMin()) <= 0);
+    const ShardId& shardId,
+    const ChunkRange& chunkRange,
+    const ChunkVersion& version) {
 
     auto commandInfo = std::make_shared<MergeChunksCommandInfo>(
-        nss,
-        lowerBound.getShard(),
-        lowerBound.getMin(),
-        upperBound.getMax(),
-        lowerBound.getVersion().isOlderThan(upperBound.getVersion()) ? upperBound.getVersion()
-                                                                     : lowerBound.getVersion());
+        nss, shardId, chunkRange.getMin(), chunkRange.getMax(), version);
 
     auto requestCollectionInfo = _buildAndEnqueueNewRequest(opCtx, std::move(commandInfo));
     return std::make_unique<MergeChunksResponseImpl>(std::move(requestCollectionInfo));
@@ -161,12 +155,12 @@ std::unique_ptr<SplitVectorResponse> BalancerCommandsSchedulerImpl::requestSplit
     OperationContext* opCtx,
     const NamespaceString& nss,
     const ChunkType& chunk,
-    const ShardKeyPattern& shardKeyPattern,
+    const KeyPattern& keyPattern,
     const SplitVectorSettings& commandSettings) {
 
     auto commandInfo = std::make_shared<SplitVectorCommandInfo>(nss,
                                                                 chunk.getShard(),
-                                                                shardKeyPattern.toBSON(),
+                                                                keyPattern.toBSON(),
                                                                 chunk.getMin(),
                                                                 chunk.getMax(),
                                                                 commandSettings.maxSplitPoints,
@@ -182,7 +176,7 @@ std::unique_ptr<SplitChunkResponse> BalancerCommandsSchedulerImpl::requestSplitC
     OperationContext* opCtx,
     const NamespaceString& nss,
     const ChunkType& chunk,
-    const ShardKeyPattern& shardKeyPattern,
+    const KeyPattern& shardKeyPattern,
     const std::vector<BSONObj>& splitPoints) {
 
     auto commandInfo = std::make_shared<SplitChunkCommandInfo>(nss,
@@ -197,19 +191,21 @@ std::unique_ptr<SplitChunkResponse> BalancerCommandsSchedulerImpl::requestSplitC
     return std::make_unique<SplitChunkResponseImpl>(std::move(requestCollectionInfo));
 }
 
-std::unique_ptr<ChunkDataSizeResponse> BalancerCommandsSchedulerImpl::requestChunkDataSize(
+std::unique_ptr<ChunkDataSizeResponse> BalancerCommandsSchedulerImpl::requestDataSize(
     OperationContext* opCtx,
     const NamespaceString& nss,
-    const ChunkType& chunk,
-    const ShardKeyPattern& shardKeyPattern,
+    const ShardId& shardId,
+    const ChunkRange& chunkRange,
+    const ChunkVersion& version,
+    const KeyPattern& keyPattern,
     bool estimatedValue) {
     auto commandInfo = std::make_shared<DataSizeCommandInfo>(nss,
-                                                             chunk.getShard(),
-                                                             shardKeyPattern.toBSON(),
-                                                             chunk.getMin(),
-                                                             chunk.getMax(),
+                                                             shardId,
+                                                             keyPattern.toBSON(),
+                                                             chunkRange.getMin(),
+                                                             chunkRange.getMax(),
                                                              estimatedValue,
-                                                             chunk.getVersion());
+                                                             version);
 
     auto requestCollectionInfo = _buildAndEnqueueNewRequest(opCtx, std::move(commandInfo));
     return std::make_unique<ChunkDataSizeResponseImpl>(std::move(requestCollectionInfo));
diff --git a/src/mongo/db/s/balancer/balancer_commands_scheduler_impl.h b/src/mongo/db/s/balancer/balancer_commands_scheduler_impl.h
index 6970d5a965c..a88f0ccf91a 100644
--- a/src/mongo/db/s/balancer/balancer_commands_scheduler_impl.h
+++ b/src/mongo/db/s/balancer/balancer_commands_scheduler_impl.h
@@ -37,6 +37,7 @@
 #include "mongo/db/s/dist_lock_manager.h"
 #include "mongo/db/service_context.h"
 #include "mongo/platform/mutex.h"
+#include "mongo/rpc/get_status_from_command_result.h"
 #include "mongo/stdx/condition_variable.h"
 #include "mongo/stdx/thread.h"
 #include "mongo/util/concurrency/notification.h"
@@ -86,7 +87,11 @@ public:
     }
 
     Status getOutcome() {
-        return getRemoteResponse().status;
+        auto response = getRemoteResponse();
+        if (!response.status.isOK()) {
+            return response.status;
+        }
+        return getStatusFromCommandResult(response.data);
     }
 
     executor::RemoteCommandResponse getRemoteResponse() {
@@ -165,6 +170,11 @@ public:
         if (!response.status.isOK()) {
             return response.status;
         }
+        auto commandStatus = getStatusFromCommandResult(response.data);
+        if (!commandStatus.isOK()) {
+            return commandStatus;
+        }
+
         std::vector<BSONObj> splitKeys;
         BSONObjIterator it(response.data.getObjectField("splitKeys"));
         while (it.more()) {
@@ -214,6 +224,10 @@ public:
         if (!response.status.isOK()) {
             return response.status;
         }
+        auto commandStatus = getStatusFromCommandResult(response.data);
+        if (!commandStatus.isOK()) {
+            return commandStatus;
+        }
         return response.data["size"].number();
     }
 
@@ -222,6 +236,10 @@ public:
         if (!response.status.isOK()) {
             return response.status;
         }
+        auto commandStatus = getStatusFromCommandResult(response.data);
+        if (!commandStatus.isOK()) {
+            return commandStatus;
+        }
         return response.data["numObjects"].number();
     }
 };
@@ -648,29 +666,31 @@ public:
 
     std::unique_ptr<MergeChunksResponse> requestMergeChunks(OperationContext* opCtx,
                                                             const NamespaceString& nss,
-                                                            const ChunkType& lowerBound,
-                                                            const ChunkType& upperBound) override;
+                                                            const ShardId& shardId,
+                                                            const ChunkRange& chunkRange,
+                                                            const ChunkVersion& version) override;
 
     std::unique_ptr<SplitVectorResponse> requestSplitVector(
         OperationContext* opCtx,
         const NamespaceString& nss,
         const ChunkType& chunk,
-        const ShardKeyPattern& shardKeyPattern,
+        const KeyPattern& keyPattern,
         const SplitVectorSettings& commandSettings) override;
 
     std::unique_ptr<SplitChunkResponse> requestSplitChunk(
         OperationContext* opCtx,
         const NamespaceString& nss,
         const ChunkType& chunk,
-        const ShardKeyPattern& shardKeyPattern,
+        const KeyPattern& keyPattern,
         const std::vector<BSONObj>& splitPoints) override;
 
-    std::unique_ptr<ChunkDataSizeResponse> requestChunkDataSize(
-        OperationContext* opCtx,
-        const NamespaceString& nss,
-        const ChunkType& chunk,
-        const ShardKeyPattern& shardKeyPattern,
-        bool estimatedValue) override;
+    std::unique_ptr<ChunkDataSizeResponse> requestDataSize(OperationContext* opCtx,
+                                                           const NamespaceString& nss,
+                                                           const ShardId& shardId,
+                                                           const ChunkRange& chunkRange,
+                                                           const ChunkVersion& version,
+                                                           const KeyPattern& keyPattern,
+                                                           bool estimatedValue) override;
 
 private:
     enum class SchedulerState { Recovering, Running, Stopping, Stopped };
diff --git a/src/mongo/db/s/balancer/balancer_commands_scheduler_test.cpp b/src/mongo/db/s/balancer/balancer_commands_scheduler_test.cpp
index e449c829eaa..41db6025b56 100644
--- a/src/mongo/db/s/balancer/balancer_commands_scheduler_test.cpp
+++ b/src/mongo/db/s/balancer/balancer_commands_scheduler_test.cpp
@@ -147,13 +147,14 @@ TEST_F(BalancerCommandsSchedulerTest, SuccessfulMoveChunkCommand) {
 
 TEST_F(BalancerCommandsSchedulerTest, SuccessfulMergeChunkCommand) {
     _scheduler.start(operationContext());
-    ChunkType chunk1 = makeChunk(0, kShardId0);
-    ChunkType chunk2 = makeChunk(10, kShardId0);
     auto networkResponseFuture = launchAsync([&]() {
         onCommand(
             [&](const executor::RemoteCommandRequest& request) { return BSON("ok" << true); });
     });
-    auto resp = _scheduler.requestMergeChunks(operationContext(), kNss, chunk1, chunk2);
+
+    ChunkRange range(BSON("x" << 0), BSON("x" << 20));
+    ChunkVersion version(1, 1, OID::gen(), Timestamp(10));
+    auto resp = _scheduler.requestMergeChunks(operationContext(), kNss, kShardId0, range, version);
     ASSERT_OK(resp->getOutcome());
     networkResponseFuture.default_timed_get();
     _scheduler.stop();
@@ -161,11 +162,10 @@ TEST_F(BalancerCommandsSchedulerTest, SuccessfulMergeChunkCommand) {
 
 TEST_F(BalancerCommandsSchedulerTest, MergeChunkNonexistentShard) {
     _scheduler.start(operationContext());
-    ChunkType brokenChunk1 = makeChunk(0, kShardId0);
-    brokenChunk1.setShard(ShardId("nonexistent"));
-    ChunkType brokenChunk2 = makeChunk(10, kShardId0);
-    brokenChunk2.setShard(ShardId("nonexistent"));
-    auto resp = _scheduler.requestMergeChunks(operationContext(), kNss, brokenChunk1, brokenChunk2);
+    ChunkRange range(BSON("x" << 0), BSON("x" << 20));
+    ChunkVersion version(1, 1, OID::gen(), Timestamp(10));
+    auto resp = _scheduler.requestMergeChunks(
+        operationContext(), kNss, ShardId("nonexistent"), range, version);
     auto shardNotFoundError = Status{ErrorCodes::ShardNotFound, "Shard nonexistent not found"};
     ASSERT_EQ(resp->getOutcome(), shardNotFoundError);
     _scheduler.stop();
@@ -184,11 +184,8 @@ TEST_F(BalancerCommandsSchedulerTest, SuccessfulSplitVectorCommand) {
             return splitChunkResponse.obj();
         });
     });
-    auto resp = _scheduler.requestSplitVector(operationContext(),
-                                              kNss,
-                                              splitChunk,
-                                              ShardKeyPattern(BSON("x" << 1)),
-                                              SplitVectorSettings());
+    auto resp = _scheduler.requestSplitVector(
+        operationContext(), kNss, splitChunk, KeyPattern(BSON("x" << 1)), SplitVectorSettings());
     ASSERT_OK(resp->getOutcome());
     ASSERT_OK(resp->getSplitKeys().getStatus());
     ASSERT_EQ(resp->getSplitKeys().getValue().size(), 1);
@@ -207,7 +204,7 @@ TEST_F(BalancerCommandsSchedulerTest, SuccessfulSplitChunkCommand) {
     auto resp = _scheduler.requestSplitChunk(operationContext(),
                                              kNss,
                                              splitChunk,
-                                             ShardKeyPattern(BSON("x" << 1)),
+                                             KeyPattern(BSON("x" << 1)),
                                              std::vector<BSONObj>{BSON("x" << 5)});
     ASSERT_OK(resp->getOutcome());
     networkResponseFuture.default_timed_get();
@@ -225,8 +222,13 @@ TEST_F(BalancerCommandsSchedulerTest, SuccessfulRequestChunkDataSizeCommand) {
         onCommand(
             [&](const executor::RemoteCommandRequest& request) { return chunkSizeResponse.obj(); });
     });
-    auto resp = _scheduler.requestChunkDataSize(
-        operationContext(), kNss, chunk, ShardKeyPattern(BSON("x" << 1)), false);
+    auto resp = _scheduler.requestDataSize(operationContext(),
+                                           kNss,
+                                           chunk.getShard(),
+                                           chunk.getRange(),
+                                           chunk.getVersion(),
+                                           KeyPattern(BSON("x" << 1)),
+                                           false);
     ASSERT_OK(resp->getOutcome());
     ASSERT_OK(resp->getSize().getStatus());
     ASSERT_EQ(resp->getSize().getValue(), 156);