SERVER-26631 Implement multiple document support for StorageInterface::findDocuments/deleteDocuments

3 files changed, 207 insertions, 51 deletions

diff --git a/src/mongo/db/repl/oplog_buffer_collection.cpp b/src/mongo/db/repl/oplog_buffer_collection.cpp
index 6e1be76806b..619c7131169 100644
--- a/src/mongo/db/repl/oplog_buffer_collection.cpp
+++ b/src/mongo/db/repl/oplog_buffer_collection.cpp
@@ -232,6 +232,8 @@ bool OplogBufferCollection::_pop_inlock(OperationContext* txn, Value* value) {
 bool OplogBufferCollection::_peekOneSide_inlock(OperationContext* txn,
                                                 Value* value,
                                                 bool front) const {
+    invariant(_count > 0);
+
     // If there is a sentinel, and it was pushed right after the last BSONObj to be popped was
     // pushed, then we return an empty BSONObj for the sentinel.
     if (!_sentinels.empty() && (_lastPoppedTimestamp == _sentinels.front())) {
@@ -251,16 +253,11 @@ bool OplogBufferCollection::_peekOneSide_inlock(OperationContext* txn,
         boundInclusion = BoundInclusion::kIncludeEndKeyOnly;
     }
 
-    auto result = _storageInterface->findDocuments(
-        txn, _nss, kIdIdxName, scanDirection, startKey, boundInclusion, 1U);
-    if (!result.isOK()) {
-        if (result != ErrorCodes::CollectionIsEmpty) {
-            fassert(40163, result.getStatus());
-        }
-        return false;
-    }
-    auto&& docs = result.getValue();
-    invariant(!docs.empty());
+    const auto docs =
+        fassertStatusOK(40163,
+                        _storageInterface->findDocuments(
+                            txn, _nss, kIdIdxName, scanDirection, startKey, boundInclusion, 1U));
+    invariant(1U == docs.size());
     *value = extractEmbeddedOplogDocument(docs.front()).getOwned();
     return true;
 }
diff --git a/src/mongo/db/repl/storage_interface_impl.cpp b/src/mongo/db/repl/storage_interface_impl.cpp
index 2ccb8006282..ce03ab1629c 100644
--- a/src/mongo/db/repl/storage_interface_impl.cpp
+++ b/src/mongo/db/repl/storage_interface_impl.cpp
@@ -443,9 +443,9 @@ namespace {
 /**
  * Returns DeleteStageParams for deleteOne with fetch.
  */
-DeleteStageParams makeDeleteStageParamsForDeleteOne() {
+DeleteStageParams makeDeleteStageParamsForDeleteDocuments() {
     DeleteStageParams deleteStageParams;
-    invariant(!deleteStageParams.isMulti);
+    deleteStageParams.isMulti = true;
     deleteStageParams.returnDeleted = true;
     return deleteStageParams;
 }
@@ -463,8 +463,6 @@ StatusWith<std::vector<BSONObj>> _findOrDeleteDocuments(
     BoundInclusion boundInclusion,
     std::size_t limit,
     FindDeleteMode mode) {
-    invariant(limit == 1U);
-
     auto isFind = mode == FindDeleteMode::kFind;
     auto opStr = isFind ? "StorageInterfaceImpl::findOne" : "StorageInterfaceImpl::deleteOne";
 
@@ -496,11 +494,12 @@ StatusWith<std::vector<BSONObj>> _findOrDeleteDocuments(
             planExecutor = isFind
                 ? InternalPlanner::collectionScan(
                       txn, nss.ns(), collection, PlanExecutor::YIELD_MANUAL, direction)
-                : InternalPlanner::deleteWithCollectionScan(txn,
-                                                            collection,
-                                                            makeDeleteStageParamsForDeleteOne(),
-                                                            PlanExecutor::YIELD_MANUAL,
-                                                            direction);
+                : InternalPlanner::deleteWithCollectionScan(
+                      txn,
+                      collection,
+                      makeDeleteStageParamsForDeleteDocuments(),
+                      PlanExecutor::YIELD_MANUAL,
+                      direction);
         } else {
             // Use index scan.
             auto indexCatalog = collection->getIndexCatalog();
@@ -541,7 +540,7 @@ StatusWith<std::vector<BSONObj>> _findOrDeleteDocuments(
                                              InternalPlanner::IXSCAN_FETCH)
                 : InternalPlanner::deleteWithIndexScan(txn,
                                                        collection,
-                                                       makeDeleteStageParamsForDeleteOne(),
+                                                       makeDeleteStageParamsForDeleteDocuments(),
                                                        indexDescriptor,
                                                        bounds.first,
                                                        bounds.second,
@@ -550,15 +549,17 @@ StatusWith<std::vector<BSONObj>> _findOrDeleteDocuments(
                                                        direction);
         }
 
-        BSONObj doc;
-        auto state = planExecutor->getNext(&doc, nullptr);
-        if (PlanExecutor::IS_EOF == state) {
-            return {ErrorCodes::CollectionIsEmpty,
-                    str::stream() << "Collection is empty, ns: " << nss.ns()};
-        }
-        invariant(PlanExecutor::ADVANCED == state);
         std::vector<BSONObj> docs;
-        docs.push_back(doc.getOwned());
+        while (docs.size() < limit) {
+            BSONObj doc;
+            auto state = planExecutor->getNext(&doc, nullptr);
+            if (PlanExecutor::ADVANCED == state) {
+                docs.push_back(doc.getOwned());
+            } else {
+                invariant(PlanExecutor::IS_EOF == state);
+                break;
+            }
+        }
         return docs;
     }
     MONGO_WRITE_CONFLICT_RETRY_LOOP_END(txn, opStr, nss.ns());
diff --git a/src/mongo/db/repl/storage_interface_impl_test.cpp b/src/mongo/db/repl/storage_interface_impl_test.cpp
index 3382f803ebb..f808a5d2bcd 100644
--- a/src/mongo/db/repl/storage_interface_impl_test.cpp
+++ b/src/mongo/db/repl/storage_interface_impl_test.cpp
@@ -653,23 +653,36 @@ TEST_F(StorageInterfaceImplWithReplCoordTest,
                       .getStatus());
 }
 
-TEST_F(StorageInterfaceImplWithReplCoordTest,
-       FindDocumentsReturnsCollectionIsEmptyIfCollectionIsEmpty) {
+TEST_F(StorageInterfaceImplWithReplCoordTest, FindDocumentsReturnsEmptyVectorIfCollectionIsEmpty) {
     auto txn = getOperationContext();
     StorageInterfaceImpl storage;
     auto nss = makeNamespace(_agent);
     auto indexName = "_id_"_sd;
     ASSERT_OK(storage.createCollection(txn, nss, CollectionOptions()));
-    ASSERT_EQUALS(ErrorCodes::CollectionIsEmpty,
-                  storage
-                      .findDocuments(txn,
-                                     nss,
-                                     indexName,
-                                     StorageInterface::ScanDirection::kForward,
-                                     {},
-                                     BoundInclusion::kIncludeStartKeyOnly,
-                                     1U)
-                      .getStatus());
+    ASSERT_TRUE(unittest::assertGet(storage.findDocuments(txn,
+                                                          nss,
+                                                          indexName,
+                                                          StorageInterface::ScanDirection::kForward,
+                                                          {},
+                                                          BoundInclusion::kIncludeStartKeyOnly,
+                                                          1U))
+                    .empty());
+}
+
+std::string _toString(const std::vector<BSONObj>& docs) {
+    str::stream ss;
+    ss << "[";
+    bool first = true;
+    for (const auto& doc : docs) {
+        if (first) {
+            ss << doc;
+            first = false;
+        } else {
+            ss << ", " << doc;
+        }
+    }
+    ss << "]";
+    return ss;
 }
 
 /**
@@ -689,6 +702,22 @@ void _assertDocumentsInCollectionEquals(OperationContext* txn,
 }
 
 /**
+ * Check StatusWith<std::vector<BSONObj>> value.
+ */
+void _assertDocumentsEqual(const StatusWith<std::vector<BSONObj>>& statusWithDocs,
+                           const std::vector<BSONObj>& expectedDocs) {
+    const auto actualDocs = unittest::assertGet(statusWithDocs);
+    auto iter = actualDocs.cbegin();
+    std::string msg = str::stream() << "expected: " << _toString(expectedDocs)
+                                    << "; actual: " << _toString(actualDocs);
+    for (const auto& doc : expectedDocs) {
+        ASSERT_TRUE(iter != actualDocs.cend()) << msg;
+        ASSERT_BSONOBJ_EQ(doc, *(iter++));
+    }
+    ASSERT_TRUE(iter == actualDocs.cend()) << msg;
+}
+
+/**
  * Returns first BSONObj from a StatusWith<std::vector<BSONObj>>.
  */
 BSONObj _assetGetFront(const StatusWith<std::vector<BSONObj>>& statusWithDocs) {
@@ -723,6 +752,26 @@ TEST_F(StorageInterfaceImplWithReplCoordTest,
                                              BoundInclusion::kIncludeStartKeyOnly,
                                              1U)));
 
+    // startKey not provided. limit is 0.
+    _assertDocumentsEqual(storage.findDocuments(txn,
+                                                nss,
+                                                indexName,
+                                                StorageInterface::ScanDirection::kForward,
+                                                {},
+                                                BoundInclusion::kIncludeStartKeyOnly,
+                                                0U),
+                          {});
+
+    // startKey not provided. limit of 2.
+    _assertDocumentsEqual(storage.findDocuments(txn,
+                                                nss,
+                                                indexName,
+                                                StorageInterface::ScanDirection::kForward,
+                                                {},
+                                                BoundInclusion::kIncludeStartKeyOnly,
+                                                2U),
+                          {BSON("_id" << 0), BSON("_id" << 1)});
+
     // startKey provided; include start key
     ASSERT_BSONOBJ_EQ(
         BSON("_id" << 0),
@@ -796,6 +845,17 @@ TEST_F(StorageInterfaceImplWithReplCoordTest,
                                              BoundInclusion::kExcludeBothStartAndEndKeys,
                                              1U)));
 
+    // startKey provided; exclude both start and end keys.
+    // A limit of 3 should return 2 documents because we reached the end of the collection.
+    _assertDocumentsEqual(storage.findDocuments(txn,
+                                                nss,
+                                                indexName,
+                                                StorageInterface::ScanDirection::kForward,
+                                                BSON("" << 2),
+                                                BoundInclusion::kExcludeBothStartAndEndKeys,
+                                                3U),
+                          {BSON("_id" << 3), BSON("_id" << 4)});
+
     _assertDocumentsInCollectionEquals(
         txn,
         nss,
@@ -828,6 +888,26 @@ TEST_F(StorageInterfaceImplWithReplCoordTest,
                                              BoundInclusion::kIncludeStartKeyOnly,
                                              1U)));
 
+    // startKey not provided. limit is 0.
+    _assertDocumentsEqual(storage.findDocuments(txn,
+                                                nss,
+                                                indexName,
+                                                StorageInterface::ScanDirection::kBackward,
+                                                {},
+                                                BoundInclusion::kIncludeStartKeyOnly,
+                                                0U),
+                          {});
+
+    // startKey not provided. limit of 2.
+    _assertDocumentsEqual(storage.findDocuments(txn,
+                                                nss,
+                                                indexName,
+                                                StorageInterface::ScanDirection::kBackward,
+                                                {},
+                                                BoundInclusion::kIncludeStartKeyOnly,
+                                                2U),
+                          {BSON("_id" << 4), BSON("_id" << 3)});
+
     // startKey provided; include start key
     ASSERT_BSONOBJ_EQ(
         BSON("_id" << 4),
@@ -881,6 +961,17 @@ TEST_F(StorageInterfaceImplWithReplCoordTest,
                                              BoundInclusion::kExcludeBothStartAndEndKeys,
                                              1U)));
 
+    // startKey provided; exclude both start and end keys.
+    // A limit of 3 should return 2 documents because we reached the beginning of the collection.
+    _assertDocumentsEqual(storage.findDocuments(txn,
+                                                nss,
+                                                indexName,
+                                                StorageInterface::ScanDirection::kBackward,
+                                                BSON("" << 2),
+                                                BoundInclusion::kExcludeBothStartAndEndKeys,
+                                                3U),
+                          {BSON("_id" << 1), BSON("_id" << 0)});
+
     _assertDocumentsInCollectionEquals(
         txn,
         nss,
@@ -1013,22 +1104,21 @@ TEST_F(StorageInterfaceImplWithReplCoordTest, DeleteDocumentsReturnsIndexNotFoun
 }
 
 TEST_F(StorageInterfaceImplWithReplCoordTest,
-       DeleteDocumentsReturnsCollectionIsEmptyIfCollectionIsEmpty) {
+       DeleteDocumentsReturnsEmptyVectorIfCollectionIsEmpty) {
     auto txn = getOperationContext();
     StorageInterfaceImpl storage;
     auto nss = makeNamespace(_agent);
     auto indexName = "_id_"_sd;
     ASSERT_OK(storage.createCollection(txn, nss, CollectionOptions()));
-    ASSERT_EQUALS(ErrorCodes::CollectionIsEmpty,
-                  storage
-                      .deleteDocuments(txn,
-                                       nss,
-                                       indexName,
-                                       StorageInterface::ScanDirection::kForward,
-                                       {},
-                                       BoundInclusion::kIncludeStartKeyOnly,
-                                       1U)
-                      .getStatus());
+    ASSERT_TRUE(
+        unittest::assertGet(storage.deleteDocuments(txn,
+                                                    nss,
+                                                    indexName,
+                                                    StorageInterface::ScanDirection::kForward,
+                                                    {},
+                                                    BoundInclusion::kIncludeStartKeyOnly,
+                                                    1U))
+            .empty());
 }
 
 TEST_F(StorageInterfaceImplWithReplCoordTest,
@@ -1070,6 +1160,26 @@ TEST_F(StorageInterfaceImplWithReplCoordTest,
                                         BSON("_id" << 6),
                                         BSON("_id" << 7)});
 
+    // startKey not provided. limit is 0.
+    _assertDocumentsEqual(storage.deleteDocuments(txn,
+                                                  nss,
+                                                  indexName,
+                                                  StorageInterface::ScanDirection::kForward,
+                                                  {},
+                                                  BoundInclusion::kIncludeStartKeyOnly,
+                                                  0U),
+                          {});
+
+    _assertDocumentsInCollectionEquals(txn,
+                                       nss,
+                                       {BSON("_id" << 1),
+                                        BSON("_id" << 2),
+                                        BSON("_id" << 3),
+                                        BSON("_id" << 4),
+                                        BSON("_id" << 5),
+                                        BSON("_id" << 6),
+                                        BSON("_id" << 7)});
+
     // startKey provided; include start key
     ASSERT_BSONOBJ_EQ(
         BSON("_id" << 2),
@@ -1105,6 +1215,20 @@ TEST_F(StorageInterfaceImplWithReplCoordTest,
         txn,
         nss,
         {BSON("_id" << 1), BSON("_id" << 3), BSON("_id" << 4), BSON("_id" << 6), BSON("_id" << 7)});
+
+    // startKey provided; exclude start key.
+    // A limit of 3 should return 2 documents because we reached the end of the collection.
+    _assertDocumentsEqual(storage.deleteDocuments(txn,
+                                                  nss,
+                                                  indexName,
+                                                  StorageInterface::ScanDirection::kForward,
+                                                  BSON("" << 4),
+                                                  BoundInclusion::kIncludeEndKeyOnly,
+                                                  3U),
+                          {BSON("_id" << 6), BSON("_id" << 7)});
+
+    _assertDocumentsInCollectionEquals(
+        txn, nss, {BSON("_id" << 1), BSON("_id" << 3), BSON("_id" << 4)});
 }
 
 TEST_F(StorageInterfaceImplWithReplCoordTest,
@@ -1146,6 +1270,26 @@ TEST_F(StorageInterfaceImplWithReplCoordTest,
                                         BSON("_id" << 5),
                                         BSON("_id" << 6)});
 
+    // startKey not provided. limit is 0.
+    _assertDocumentsEqual(storage.deleteDocuments(txn,
+                                                  nss,
+                                                  indexName,
+                                                  StorageInterface::ScanDirection::kBackward,
+                                                  {},
+                                                  BoundInclusion::kIncludeStartKeyOnly,
+                                                  0U),
+                          {});
+
+    _assertDocumentsInCollectionEquals(txn,
+                                       nss,
+                                       {BSON("_id" << 0),
+                                        BSON("_id" << 1),
+                                        BSON("_id" << 2),
+                                        BSON("_id" << 3),
+                                        BSON("_id" << 4),
+                                        BSON("_id" << 5),
+                                        BSON("_id" << 6)});
+
     // startKey provided; include start key
     ASSERT_BSONOBJ_EQ(
         BSON("_id" << 5),
@@ -1181,6 +1325,20 @@ TEST_F(StorageInterfaceImplWithReplCoordTest,
         txn,
         nss,
         {BSON("_id" << 0), BSON("_id" << 1), BSON("_id" << 3), BSON("_id" << 4), BSON("_id" << 6)});
+
+    // startKey provided; exclude start key.
+    // A limit of 3 should return 2 documents because we reached the beginning of the collection.
+    _assertDocumentsEqual(storage.deleteDocuments(txn,
+                                                  nss,
+                                                  indexName,
+                                                  StorageInterface::ScanDirection::kBackward,
+                                                  BSON("" << 3),
+                                                  BoundInclusion::kIncludeEndKeyOnly,
+                                                  3U),
+                          {BSON("_id" << 1), BSON("_id" << 0)});
+
+    _assertDocumentsInCollectionEquals(
+        txn, nss, {BSON("_id" << 3), BSON("_id" << 4), BSON("_id" << 6)});
 }
 
 TEST_F(StorageInterfaceImplWithReplCoordTest,