SERVER-61678 Make autoSplitVector return `continuation` flag

17 files changed, 132 insertions, 111 deletions

diff --git a/src/mongo/db/s/auto_split_vector.cpp b/src/mongo/db/s/auto_split_vector.cpp
index a8f30e5742c..c343cbed7f2 100644
--- a/src/mongo/db/s/auto_split_vector.cpp
+++ b/src/mongo/db/s/auto_split_vector.cpp
@@ -90,13 +90,14 @@ auto orderShardKeyFields(const BSONObj& keyPattern, BSONObj& key) {
 
 }  // namespace
 
-std::vector<BSONObj> autoSplitVector(OperationContext* opCtx,
-                                     const NamespaceString& nss,
-                                     const BSONObj& keyPattern,
-                                     const BSONObj& min,
-                                     const BSONObj& max,
-                                     long long maxChunkSizeBytes) {
+std::pair<std::vector<BSONObj>, bool> autoSplitVector(OperationContext* opCtx,
+                                                      const NamespaceString& nss,
+                                                      const BSONObj& keyPattern,
+                                                      const BSONObj& min,
+                                                      const BSONObj& max,
+                                                      long long maxChunkSizeBytes) {
     std::vector<BSONObj> splitKeys;
+    bool reachedMaxBSONSize = false;  // True if the split points vector becomes too big
 
     int elapsedMillisToFindSplitPoints;
 
@@ -201,7 +202,6 @@ std::vector<BSONObj> autoSplitVector(OperationContext* opCtx,
         BSONObj currentKey;               // Last key seen during the index scan
         long long numScannedKeys = 1;     // minKeyInOriginalChunk has already been scanned
         std::size_t resultArraySize = 0;  // Approximate size in bytes of the split points array
-        bool reachedMaxBSONSize = false;  // True if the split points vector becomes too big
 
         // Lambda to check whether the split points vector would exceed BSONObjMaxUserSize in case
         // of additional split key of the specified size.
@@ -351,7 +351,7 @@ std::vector<BSONObj> autoSplitVector(OperationContext* opCtx,
                       "duration"_attr = Milliseconds(elapsedMillisToFindSplitPoints));
     }
 
-    return splitKeys;
+    return std::make_pair(splitKeys, reachedMaxBSONSize);
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/s/auto_split_vector.h b/src/mongo/db/s/auto_split_vector.h
index 559c1f814d6..d64b525d59e 100644
--- a/src/mongo/db/s/auto_split_vector.h
+++ b/src/mongo/db/s/auto_split_vector.h
@@ -36,7 +36,10 @@ namespace mongo {
 
 /**
  * Given a chunk, determines whether it satisfies the requisites to be auto-splitted and - if so -
- * returns the split points (shard keys representing the lower bounds of the new chunks to create).
+ * returns the split points (shard keys representing the lower bounds of the new chunks to create)
+ * along with a bool representing whether or not the max bson size for the vector of split points
+ * was reached. If the bool returns true, this function should be called on the same range again to
+ * retrieve the rest of the split points.
  *
  * The logic implemented can be summarized as follows: given a `maxChunkSize` of `x` MB, the
  * algorithm aims to choose the split points so that the resulting chunks' size would be around
@@ -74,6 +77,7 @@ namespace mongo {
  * (2.1) `S >= 80% maxChunkSize`, so keep the current split points.
  *
  * Returned split points: [99].
+ * Returned continuation flag: false.
  *
  * ========= EXAMPLE (CASE 2.2) ========
  * `maxChunkSize` = 100MB
@@ -89,6 +93,7 @@ namespace mongo {
  * 67% maxChunkSize. Recalculate.
  *
  * Returned split points: [69].
+ * Returned continuation flag: false.
  *
  * ========= EXAMPLE (CASE 2.3) ========
  * `maxChunkSize` = 100MB
@@ -104,13 +109,14 @@ namespace mongo {
  * >= 67% maxChunkSize. So remove the last split point.
  *
  * Returned split points: [].
+ * Returned continuation flag: false.
  */
-std::vector<BSONObj> autoSplitVector(OperationContext* opCtx,
-                                     const NamespaceString& nss,
-                                     const BSONObj& keyPattern,
-                                     const BSONObj& min,
-                                     const BSONObj& max,
-                                     long long maxChunkSizeBytes);
+std::pair<std::vector<BSONObj>, bool> autoSplitVector(OperationContext* opCtx,
+                                                      const NamespaceString& nss,
+                                                      const BSONObj& keyPattern,
+                                                      const BSONObj& min,
+                                                      const BSONObj& max,
+                                                      long long maxChunkSizeBytes);
 
 /*
  * Utility function for deserializing autoSplitVector/splitVector responses.
diff --git a/src/mongo/db/s/auto_split_vector_command.cpp b/src/mongo/db/s/auto_split_vector_command.cpp
index 1748491cac6..1681c011f97 100644
--- a/src/mongo/db/s/auto_split_vector_command.cpp
+++ b/src/mongo/db/s/auto_split_vector_command.cpp
@@ -79,14 +79,15 @@ public:
                                   << kSmallestChunkSizeSupported,
                     req.getMaxChunkSizeBytes() >= kSmallestChunkSizeSupported);
 
-            auto splitKeys = autoSplitVector(opCtx,
-                                             ns(),
-                                             req.getKeyPattern(),
-                                             req.getMin(),
-                                             req.getMax(),
-                                             req.getMaxChunkSizeBytes());
-
-            return splitKeys;
+            auto [splitPoints, continuation] = autoSplitVector(opCtx,
+                                                               ns(),
+                                                               req.getKeyPattern(),
+                                                               req.getMin(),
+                                                               req.getMax(),
+                                                               req.getMaxChunkSizeBytes());
+            Response autoSplitVectorResponse(splitPoints);
+            autoSplitVectorResponse.setContinuation(continuation);
+            return autoSplitVectorResponse;
         }
 
     private:
diff --git a/src/mongo/db/s/auto_split_vector_test.cpp b/src/mongo/db/s/auto_split_vector_test.cpp
index 45371606529..68e99802a4e 100644
--- a/src/mongo/db/s/auto_split_vector_test.cpp
+++ b/src/mongo/db/s/auto_split_vector_test.cpp
@@ -50,7 +50,7 @@ const std::string kPattern = "_id";
  * Call the autoSplitVector function of the test collection on a chunk with bounds [0, 100) and with
  * the specified `maxChunkSizeMB`.
  */
-std::vector<BSONObj> autoSplit(OperationContext* opCtx, int maxChunkSizeMB) {
+std::pair<std::vector<BSONObj>, bool> autoSplit(OperationContext* opCtx, int maxChunkSizeMB) {
     return autoSplitVector(opCtx,
                            kNss,
                            BSON(kPattern << 1) /* shard key pattern */,
@@ -141,36 +141,42 @@ TEST_F(AutoSplitVectorTest, NoCollection) {
 }
 
 TEST_F(AutoSplitVectorTest, EmptyCollection) {
-    const auto splitKey = autoSplitVector(operationContext(),
-                                          kNss,
-                                          BSON(kPattern << 1) /* shard key pattern */,
-                                          BSON(kPattern << kMinBSONKey) /* min */,
-                                          BSON(kPattern << kMaxBSONKey) /* max */,
-                                          1 * 1024 * 1024 /* max chunk size in bytes*/);
+    const auto [splitKey, continuation] =
+        autoSplitVector(operationContext(),
+                        kNss,
+                        BSON(kPattern << 1) /* shard key pattern */,
+                        BSON(kPattern << kMinBSONKey) /* min */,
+                        BSON(kPattern << kMaxBSONKey) /* max */,
+                        1 * 1024 * 1024 /* max chunk size in bytes*/);
     ASSERT_EQ(0, splitKey.size());
+    ASSERT_FALSE(continuation);
 }
 
 TEST_F(AutoSplitVectorTest10MB, EmptyRange) {
-    const auto splitKey = autoSplitVector(operationContext(),
-                                          kNss,
-                                          BSON(kPattern << 1) /* shard key pattern */,
-                                          BSON(kPattern << kMinBSONKey) /* min */,
-                                          BSON(kPattern << -10) /* max */,
-                                          1 * 1024 * 1024 /* max chunk size in bytes*/);
+    const auto [splitKey, continuation] =
+        autoSplitVector(operationContext(),
+                        kNss,
+                        BSON(kPattern << 1) /* shard key pattern */,
+                        BSON(kPattern << kMinBSONKey) /* min */,
+                        BSON(kPattern << -10) /* max */,
+                        1 * 1024 * 1024 /* max chunk size in bytes*/);
     ASSERT_EQ(0, splitKey.size());
+    ASSERT_FALSE(continuation);
 }
 
 
 // No split points if estimated `data size < max chunk size`
 TEST_F(AutoSplitVectorTest10MB, NoSplitIfDataLessThanMaxChunkSize) {
-    std::vector<BSONObj> splitKeys = autoSplit(operationContext(), 11 /* maxChunkSizeMB */);
+    auto [splitKeys, continuation] = autoSplit(operationContext(), 11 /* maxChunkSizeMB */);
     ASSERT_EQ(splitKeys.size(), 0);
+    ASSERT_FALSE(continuation);
 }
 
 // Do not split in case of `chunk size == maxChunkSize`
 TEST_F(AutoSplitVectorTest10MB, NoSplitIfDataEqualMaxChunkSize) {
-    std::vector<BSONObj> splitKeys = autoSplit(operationContext(), 10 /* maxChunkSizeMB */);
+    auto [splitKeys, continuation] = autoSplit(operationContext(), 10 /* maxChunkSizeMB */);
     ASSERT_EQ(splitKeys.size(), 0);
+    ASSERT_FALSE(continuation);
 }
 
 // No split points if `chunk size > max chunk size` but threshold not reached
@@ -181,36 +187,40 @@ TEST_F(AutoSplitVectorTest10MB, NoSplitIfDataLessThanThreshold) {
         // different range to don't interfere with the chunk getting splitted.
         insertNDocsOf1MB(operationContext(), surplus /* nDocs */);
     }
-    std::vector<BSONObj> splitKeys = autoSplit(operationContext(), 10 /* maxChunkSizeMB */);
+    auto [splitKeys, continuation] = autoSplit(operationContext(), 10 /* maxChunkSizeMB */);
     ASSERT_EQ(splitKeys.size(), 0);
+    ASSERT_FALSE(continuation);
 }
 
 // One split point if `chunk size > max chunk size` and threshold reached
 TEST_F(AutoSplitVectorTest10MB, SplitIfDataSlightlyMoreThanThreshold) {
     const auto surplus = 4;
     insertNDocsOf1MB(operationContext(), surplus /* nDocs */);
-    std::vector<BSONObj> splitKeys = autoSplit(operationContext(), 10 /* maxChunkSizeMB */);
+    auto [splitKeys, continuation] = autoSplit(operationContext(), 10 /* maxChunkSizeMB */);
     ASSERT_EQ(splitKeys.size(), 1);
     ASSERT_EQ(6, splitKeys.front().getIntField(kPattern));
+    ASSERT_FALSE(continuation);
 }
 
 // Split points if `data size > max chunk size * 2` and threshold reached
 TEST_F(AutoSplitVectorTest10MB, SplitIfDataMoreThanThreshold) {
     const auto surplus = 14;
     insertNDocsOf1MB(operationContext(), surplus /* nDocs */);
-    std::vector<BSONObj> splitKeys = autoSplit(operationContext(), 10 /* maxChunkSizeMB */);
+    auto [splitKeys, continuation] = autoSplit(operationContext(), 10 /* maxChunkSizeMB */);
     ASSERT_EQ(splitKeys.size(), 2);
     ASSERT_EQ(7, splitKeys.front().getIntField(kPattern));
     ASSERT_EQ(15, splitKeys.back().getIntField(kPattern));
+    ASSERT_FALSE(continuation);
 }
 
 // Split points are not recalculated if the right-most chunk is at least `80% maxChunkSize`
 TEST_F(AutoSplitVectorTest10MB, NoRecalculateIfBigLastChunk) {
     const auto surplus = 8;
     insertNDocsOf1MB(operationContext(), surplus /* nDocs */);
-    std::vector<BSONObj> splitKeys = autoSplit(operationContext(), 10 /* maxChunkSizeMB */);
+    auto [splitKeys, continuation] = autoSplit(operationContext(), 10 /* maxChunkSizeMB */);
     ASSERT_EQ(splitKeys.size(), 1);
     ASSERT_EQ(9, splitKeys.front().getIntField(kPattern));
+    ASSERT_FALSE(continuation);
 }
 
 class RepositionLastSplitPointsTest : public AutoSplitVectorTest {
@@ -232,7 +242,7 @@ public:
 
         int expectedChunkSize =
             mustReposition ? getInsertedSize() / (nSplitPoints + 1) : maxDocsPerChunk;
-        std::vector<BSONObj> splitKeys =
+        auto [splitKeys, continuation] =
             autoSplit(operationContext(), maxDocsPerChunk /* maxChunkSizeMB */);
 
         int approximateNextMin = expectedChunkSize;
diff --git a/src/mongo/db/s/balancer/balancer.cpp b/src/mongo/db/s/balancer/balancer.cpp
index 52bacd4d743..0271eeb5d4c 100644
--- a/src/mongo/db/s/balancer/balancer.cpp
+++ b/src/mongo/db/s/balancer/balancer.cpp
@@ -461,7 +461,7 @@ void Balancer::_consumeActionStreamLoop() {
                                                      splitVectorAction.maxChunkSizeBytes)
                             .thenRunOn(*executor)
                             .onCompletion([this, command = std::move(splitVectorAction)](
-                                              const StatusWith<std::vector<BSONObj>>&
+                                              const StatusWith<AutoSplitVectorResponse>&
                                                   swSplitPoints) {
                                 ThreadClient tc(
                                     "BalancerDefragmentationPolicy::acknowledgeSplitVectorResult",
diff --git a/src/mongo/db/s/balancer/balancer_commands_scheduler.h b/src/mongo/db/s/balancer/balancer_commands_scheduler.h
index 068ed0af167..dae383b5723 100644
--- a/src/mongo/db/s/balancer/balancer_commands_scheduler.h
+++ b/src/mongo/db/s/balancer/balancer_commands_scheduler.h
@@ -123,13 +123,14 @@ public:
                                                 const ChunkRange& chunkRange,
                                                 const ChunkVersion& version) = 0;
 
-    virtual SemiFuture<SplitPoints> requestAutoSplitVector(OperationContext* opCtx,
-                                                           const NamespaceString& nss,
-                                                           const ShardId& shardId,
-                                                           const BSONObj& keyPattern,
-                                                           const BSONObj& minKey,
-                                                           const BSONObj& maxKey,
-                                                           int64_t maxChunkSizeBytes) = 0;
+    virtual SemiFuture<AutoSplitVectorResponse> requestAutoSplitVector(
+        OperationContext* opCtx,
+        const NamespaceString& nss,
+        const ShardId& shardId,
+        const BSONObj& keyPattern,
+        const BSONObj& minKey,
+        const BSONObj& maxKey,
+        int64_t maxChunkSizeBytes) = 0;
 
     virtual SemiFuture<void> requestSplitChunk(OperationContext* opCtx,
                                                const NamespaceString& nss,
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 b8278a8f454..fe8f6561ca3 100644
--- a/src/mongo/db/s/balancer/balancer_commands_scheduler_impl.cpp
+++ b/src/mongo/db/s/balancer/balancer_commands_scheduler_impl.cpp
@@ -245,7 +245,7 @@ SemiFuture<void> BalancerCommandsSchedulerImpl::requestMergeChunks(OperationCont
         .semi();
 }
 
-SemiFuture<SplitPoints> BalancerCommandsSchedulerImpl::requestAutoSplitVector(
+SemiFuture<AutoSplitVectorResponse> BalancerCommandsSchedulerImpl::requestAutoSplitVector(
     OperationContext* opCtx,
     const NamespaceString& nss,
     const ShardId& shardId,
@@ -257,14 +257,13 @@ SemiFuture<SplitPoints> BalancerCommandsSchedulerImpl::requestAutoSplitVector(
         nss, shardId, keyPattern, minKey, maxKey, maxChunkSizeBytes);
     return _buildAndEnqueueNewRequest(opCtx, std::move(commandInfo))
         .then([](const executor::RemoteCommandResponse& remoteResponse)
-                  -> StatusWith<std::vector<BSONObj>> {
+                  -> StatusWith<AutoSplitVectorResponse> {
             auto responseStatus = processRemoteResponse(remoteResponse);
             if (!responseStatus.isOK()) {
                 return responseStatus;
             }
-            const auto payload = AutoSplitVectorResponse::parse(
-                IDLParserErrorContext("AutoSplitVectorResponse"), std::move(remoteResponse.data));
-            return payload.getSplitKeys();
+            return AutoSplitVectorResponse::parse(IDLParserErrorContext("AutoSplitVectorResponse"),
+                                                  std::move(remoteResponse.data));
         })
         .semi();
 }
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 dacc9b2bbe9..b433018d1ef 100644
--- a/src/mongo/db/s/balancer/balancer_commands_scheduler_impl.h
+++ b/src/mongo/db/s/balancer/balancer_commands_scheduler_impl.h
@@ -527,13 +527,13 @@ public:
                                         const ChunkRange& chunkRange,
                                         const ChunkVersion& version) override;
 
-    SemiFuture<SplitPoints> requestAutoSplitVector(OperationContext* opCtx,
-                                                   const NamespaceString& nss,
-                                                   const ShardId& shardId,
-                                                   const BSONObj& keyPattern,
-                                                   const BSONObj& minKey,
-                                                   const BSONObj& maxKey,
-                                                   int64_t maxChunkSizeBytes) override;
+    SemiFuture<AutoSplitVectorResponse> requestAutoSplitVector(OperationContext* opCtx,
+                                                               const NamespaceString& nss,
+                                                               const ShardId& shardId,
+                                                               const BSONObj& keyPattern,
+                                                               const BSONObj& minKey,
+                                                               const BSONObj& maxKey,
+                                                               int64_t maxChunkSizeBytes) override;
 
     SemiFuture<void> requestSplitChunk(OperationContext* opCtx,
                                        const NamespaceString& nss,
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 8ce60d3c929..4968b71ebb4 100644
--- a/src/mongo/db/s/balancer/balancer_commands_scheduler_test.cpp
+++ b/src/mongo/db/s/balancer/balancer_commands_scheduler_test.cpp
@@ -206,6 +206,7 @@ TEST_F(BalancerCommandsSchedulerTest, SuccessfulAutoSplitVectorCommand) {
     splitKeys.append(BSON("x" << 7));
     splitKeys.append(BSON("x" << 9));
     splitKeys.done();
+    autoSplitVectorResponse.append("continuation", false);
     auto networkResponseFuture = launchAsync([&]() {
         onCommand([&](const executor::RemoteCommandRequest& request) {
             return autoSplitVectorResponse.obj();
@@ -218,12 +219,14 @@ TEST_F(BalancerCommandsSchedulerTest, SuccessfulAutoSplitVectorCommand) {
                                                             splitChunk.getMin(),
                                                             splitChunk.getMax(),
                                                             4);
-    auto swReceivedSplitKeys = futureResponse.getNoThrow();
-    ASSERT_OK(swReceivedSplitKeys.getStatus());
-    auto receivedSplitKeys = swReceivedSplitKeys.getValue();
+    auto swAutoSplitVectorResponse = futureResponse.getNoThrow();
+    ASSERT_OK(swAutoSplitVectorResponse.getStatus());
+    auto receivedSplitKeys = swAutoSplitVectorResponse.getValue().getSplitKeys();
+    auto continuation = swAutoSplitVectorResponse.getValue().getContinuation();
     ASSERT_EQ(receivedSplitKeys.size(), 2);
     ASSERT_BSONOBJ_EQ(receivedSplitKeys[0], BSON("x" << 7));
     ASSERT_BSONOBJ_EQ(receivedSplitKeys[1], BSON("x" << 9));
+    ASSERT_FALSE(continuation);
     networkResponseFuture.default_timed_get();
     _scheduler.stop();
 }
diff --git a/src/mongo/db/s/balancer/balancer_defragmentation_policy.h b/src/mongo/db/s/balancer/balancer_defragmentation_policy.h
index 3cada882af9..e1a7303a5e4 100644
--- a/src/mongo/db/s/balancer/balancer_defragmentation_policy.h
+++ b/src/mongo/db/s/balancer/balancer_defragmentation_policy.h
@@ -93,9 +93,10 @@ public:
                                         MergeInfo action,
                                         const Status& result) = 0;
 
-    virtual void acknowledgeAutoSplitVectorResult(OperationContext* opCtx,
-                                                  AutoSplitVectorInfo action,
-                                                  const StatusWith<SplitPoints>& result) = 0;
+    virtual void acknowledgeAutoSplitVectorResult(
+        OperationContext* opCtx,
+        AutoSplitVectorInfo action,
+        const StatusWith<AutoSplitVectorResponse>& result) = 0;
 
     virtual void acknowledgeSplitResult(OperationContext* opCtx,
                                         SplitInfoWithKeyPattern action,
diff --git a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
index fd9c3c58754..0b1d26b0296 100644
--- a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
+++ b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
@@ -1266,12 +1266,6 @@ public:
         return boost::none;
     }
 
-    bool moreSplitPointsToReceive(const SplitPoints& splitPoints) {
-        auto addBSONSize = [](const int& size, const BSONObj& obj) { return size + obj.objsize(); };
-        int totalSize = std::accumulate(splitPoints.begin(), splitPoints.end(), 0, addBSONSize);
-        return totalSize >= BSONObjMaxUserSize - 4096;
-    }
-
     void applyActionResult(OperationContext* opCtx,
                            const DefragmentationAction& action,
                            const DefragmentationActionResponse& response) override {
@@ -1288,7 +1282,8 @@ public:
                     uasserted(ErrorCodes::BadValue, "Unexpected action type");
                 },
                 [&](const AutoSplitVectorInfo& autoSplitVectorAction) {
-                    auto& splitVectorResponse = stdx::get<StatusWith<SplitPoints>>(response);
+                    auto& splitVectorResponse =
+                        stdx::get<StatusWith<AutoSplitVectorResponse>>(response);
                     handleActionResult(
                         opCtx,
                         _nss,
@@ -1296,16 +1291,15 @@ public:
                         getType(),
                         splitVectorResponse.getStatus(),
                         [&]() {
-                            auto& splitPoints = splitVectorResponse.getValue();
+                            auto& splitPoints = splitVectorResponse.getValue().getSplitKeys();
                             if (!splitPoints.empty()) {
                                 auto& pendingActions =
                                     _pendingActionsByShards[autoSplitVectorAction.shardId];
                                 pendingActions.rangesToSplit.push_back(
                                     std::make_pair(ChunkRange(autoSplitVectorAction.minKey,
                                                               autoSplitVectorAction.maxKey),
-                                                   splitVectorResponse.getValue()));
-                                // TODO (SERVER-61678): replace with check for continuation flag
-                                if (moreSplitPointsToReceive(splitPoints)) {
+                                                   splitPoints));
+                                if (splitVectorResponse.getValue().getContinuation()) {
                                     pendingActions.rangesToFindSplitPoints.emplace_back(
                                         splitPoints.back(), autoSplitVectorAction.maxKey);
                                 }
@@ -1625,7 +1619,9 @@ void BalancerDefragmentationPolicyImpl::acknowledgeDataSizeResult(
 }
 
 void BalancerDefragmentationPolicyImpl::acknowledgeAutoSplitVectorResult(
-    OperationContext* opCtx, AutoSplitVectorInfo action, const StatusWith<SplitPoints>& result) {
+    OperationContext* opCtx,
+    AutoSplitVectorInfo action,
+    const StatusWith<AutoSplitVectorResponse>& result) {
     stdx::lock_guard<Latch> lk(_stateMutex);
     // Check if collection defragmentation has been canceled
     if (!_defragmentationStates.contains(action.uuid)) {
diff --git a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.h b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.h
index 68b30cdfc29..61fa9f86cbf 100644
--- a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.h
+++ b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.h
@@ -89,9 +89,10 @@ public:
                                 MergeInfo action,
                                 const Status& result) override;
 
-    void acknowledgeAutoSplitVectorResult(OperationContext* opCtx,
-                                          AutoSplitVectorInfo action,
-                                          const StatusWith<SplitPoints>& result) override;
+    void acknowledgeAutoSplitVectorResult(
+        OperationContext* opCtx,
+        AutoSplitVectorInfo action,
+        const StatusWith<AutoSplitVectorResponse>& result) override;
 
     void acknowledgeSplitResult(OperationContext* opCtx,
                                 SplitInfoWithKeyPattern action,
diff --git a/src/mongo/db/s/balancer/balancer_defragmentation_policy_test.cpp b/src/mongo/db/s/balancer/balancer_defragmentation_policy_test.cpp
index 05ac91fb48b..c262df03936 100644
--- a/src/mongo/db/s/balancer/balancer_defragmentation_policy_test.cpp
+++ b/src/mongo/db/s/balancer/balancer_defragmentation_policy_test.cpp
@@ -812,8 +812,10 @@ TEST_F(BalancerDefragmentationPolicyTest,
     auto autoSplitAction = stdx::get<AutoSplitVectorInfo>(future.get());
 
     std::vector<BSONObj> splitPoints{BSON("x" << 5)};
+    AutoSplitVectorResponse resp{splitPoints};
+    resp.setContinuation(false);
     _defragmentationPolicy.acknowledgeAutoSplitVectorResult(
-        operationContext(), autoSplitAction, StatusWith(splitPoints));
+        operationContext(), autoSplitAction, StatusWith(resp));
 
     // Under the setup of this test, the stream should only contain only a split action over the
     // recently AutoSplitVector-ed range.
@@ -839,8 +841,10 @@ TEST_F(BalancerDefragmentationPolicyTest, TestAutoSplitWithNoSplitPointsDoesNotT
     auto autoSplitAction = stdx::get<AutoSplitVectorInfo>(future.get());
 
     std::vector<BSONObj> splitPoints;
+    AutoSplitVectorResponse resp{splitPoints};
+    resp.setContinuation(false);
     _defragmentationPolicy.acknowledgeAutoSplitVectorResult(
-        operationContext(), autoSplitAction, StatusWith(splitPoints));
+        operationContext(), autoSplitAction, StatusWith(resp));
 
     // The stream should now be empty
     future = _defragmentationPolicy.getNextStreamingAction(operationContext());
@@ -854,21 +858,11 @@ TEST_F(BalancerDefragmentationPolicyTest, TestMoreThan16MBSplitPointsTriggersSpl
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     auto autoSplitAction = stdx::get<AutoSplitVectorInfo>(future.get());
 
-    // TODO (SERVER-61678) use continuation flag instead of large vector
-    std::vector<BSONObj> splitPoints = [] {
-        std::vector<BSONObj> splitPoints;
-        int splitPointSize = 0;
-        std::string filler(1024 * 1024, 'x');
-        int distinguisher = 0;
-        while (splitPointSize < BSONObjMaxUserSize) {
-            auto newBSON = BSON("id" << distinguisher++ << "filler" << filler);
-            splitPointSize += newBSON.objsize();
-            splitPoints.push_back(newBSON);
-        }
-        return splitPoints;
-    }();
+    std::vector<BSONObj> splitPoints{BSON("x" << 5)};
+    AutoSplitVectorResponse resp{splitPoints};
+    resp.setContinuation(true);
     _defragmentationPolicy.acknowledgeAutoSplitVectorResult(
-        operationContext(), autoSplitAction, StatusWith(splitPoints));
+        operationContext(), autoSplitAction, StatusWith(resp));
 
     // The stream should now contain one Split action with the split points from above and one
     // AutoSplitVector action from the last split point to the end of the chunk
@@ -896,8 +890,10 @@ TEST_F(BalancerDefragmentationPolicyTest, TestFailedSplitChunkActionGetsReissued
     auto autoSplitAction = stdx::get<AutoSplitVectorInfo>(future.get());
 
     std::vector<BSONObj> splitPoints{BSON("x" << 5)};
+    AutoSplitVectorResponse resp{splitPoints};
+    resp.setContinuation(false);
     _defragmentationPolicy.acknowledgeAutoSplitVectorResult(
-        operationContext(), autoSplitAction, StatusWith(splitPoints));
+        operationContext(), autoSplitAction, StatusWith(resp));
 
     // The stream should now contain the split action for the recently AutoSplitVector-ed range.
     future = _defragmentationPolicy.getNextStreamingAction(operationContext());
@@ -929,8 +925,10 @@ TEST_F(BalancerDefragmentationPolicyTest,
     auto autoSplitAction = stdx::get<AutoSplitVectorInfo>(future.get());
 
     std::vector<BSONObj> splitPoints{BSON("x" << 5)};
+    AutoSplitVectorResponse resp{splitPoints};
+    resp.setContinuation(false);
     _defragmentationPolicy.acknowledgeAutoSplitVectorResult(
-        operationContext(), autoSplitAction, StatusWith(splitPoints));
+        operationContext(), autoSplitAction, StatusWith(resp));
 
     // The stream should now contain the split action for the recently AutoSplitVector-ed range.
     future = _defragmentationPolicy.getNextStreamingAction(operationContext());
diff --git a/src/mongo/db/s/balancer/balancer_policy.h b/src/mongo/db/s/balancer/balancer_policy.h
index b86ce1cabd8..6b81068e43f 100644
--- a/src/mongo/db/s/balancer/balancer_policy.h
+++ b/src/mongo/db/s/balancer/balancer_policy.h
@@ -38,6 +38,7 @@
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/s/balancer/cluster_statistics.h"
 #include "mongo/s/catalog/type_chunk.h"
+#include "mongo/s/request_types/auto_split_vector_gen.h"
 #include "mongo/s/request_types/move_chunk_request.h"
 #include "mongo/s/shard_id.h"
 
@@ -205,7 +206,7 @@ typedef stdx::variant<MergeInfo,
                       EndOfActionStream>
     DefragmentationAction;
 
-typedef stdx::variant<Status, StatusWith<SplitPoints>, StatusWith<DataSizeResponse>>
+typedef stdx::variant<Status, StatusWith<AutoSplitVectorResponse>, StatusWith<DataSizeResponse>>
     DefragmentationActionResponse;
 
 typedef std::vector<ClusterStatistics::ShardStatistics> ShardStatisticsVector;
diff --git a/src/mongo/db/s/chunk_splitter.cpp b/src/mongo/db/s/chunk_splitter.cpp
index f26341fc66e..751e15141c1 100644
--- a/src/mongo/db/s/chunk_splitter.cpp
+++ b/src/mongo/db/s/chunk_splitter.cpp
@@ -345,12 +345,12 @@ void ChunkSplitter::_runAutosplit(std::shared_ptr<ChunkSplitStateDriver> chunkSp
                     "maxChunkSizeBytes"_attr = maxChunkSizeBytes);
 
         chunkSplitStateDriver->prepareSplit();
-        auto splitPoints = autoSplitVector(opCtx.get(),
-                                           nss,
-                                           shardKeyPattern.toBSON(),
-                                           chunk.getMin(),
-                                           chunk.getMax(),
-                                           maxChunkSizeBytes);
+        auto [splitPoints, _] = autoSplitVector(opCtx.get(),
+                                                nss,
+                                                shardKeyPattern.toBSON(),
+                                                chunk.getMin(),
+                                                chunk.getMax(),
+                                                maxChunkSizeBytes);
 
         if (splitPoints.empty()) {
             LOGV2_DEBUG(21907,
diff --git a/src/mongo/db/s/config/sharding_catalog_manager_shard_collection_test.cpp b/src/mongo/db/s/config/sharding_catalog_manager_shard_collection_test.cpp
index 06ab6ef787d..8cb835f0e6a 100644
--- a/src/mongo/db/s/config/sharding_catalog_manager_shard_collection_test.cpp
+++ b/src/mongo/db/s/config/sharding_catalog_manager_shard_collection_test.cpp
@@ -75,7 +75,7 @@ protected:
                 ReadPreferenceSetting(ReadPreference::PrimaryPreferred).toContainingBSON(),
                 rpc::TrackingMetadata::removeTrackingData(request.metadata));
 
-            return BSON("ok" << 1 << "splitKeys" << splitPoints);
+            return BSON("ok" << 1 << "splitKeys" << splitPoints << "continuation" << false);
         });
     }
 
diff --git a/src/mongo/s/request_types/auto_split_vector.idl b/src/mongo/s/request_types/auto_split_vector.idl
index d83b16766fd..abcbece2019 100644
--- a/src/mongo/s/request_types/auto_split_vector.idl
+++ b/src/mongo/s/request_types/auto_split_vector.idl
@@ -49,6 +49,10 @@ structs:
         strict: false
         fields:
             splitKeys: bson_vector
+            continuation:
+                type: bool
+                description: "Whether there are more split points available for the requested range"
+                default: false # TODO (SERVER-64350): Remove the default value after 6.0 branching.
 
 commands:
   autoSplitVector:
@@ -70,4 +74,4 @@ commands:
             description: "Max key of the chunk"
         maxChunkSizeBytes:
             type: safeInt64
-            description: "Max chunk size of the collection expressed in bytes"
\ No newline at end of file
+            description: "Max chunk size of the collection expressed in bytes"