SERVER-24365 Implement the rangesToClean set and maintenance methods

3 files changed, 200 insertions, 3 deletions

diff --git a/src/mongo/db/s/metadata_manager.cpp b/src/mongo/db/s/metadata_manager.cpp
index d62300cb4a0..9547bd64822 100644
--- a/src/mongo/db/s/metadata_manager.cpp
+++ b/src/mongo/db/s/metadata_manager.cpp
@@ -30,7 +30,6 @@
 
 #include "mongo/db/s/metadata_manager.h"
 
-#include "mongo/db/s/collection_metadata.h"
 #include "mongo/stdx/memory.h"
 
 namespace mongo {
@@ -102,4 +101,73 @@ ScopedCollectionMetadata& ScopedCollectionMetadata::operator=(ScopedCollectionMe
     return *this;
 }
 
+std::map<BSONObj, ChunkRange> MetadataManager::getCopyOfRanges() {
+    return _rangesToClean;
+}
+
+void MetadataManager::addRangeToClean(const ChunkRange& range) {
+    auto itLow = _rangesToClean.upper_bound(range.getMin());
+    if (itLow != _rangesToClean.begin()) {
+        --itLow;
+    }
+
+    if (itLow != _rangesToClean.end()) {
+        const ChunkRange& cr = itLow->second;
+        if (cr.getMin() < range.getMin()) {
+            // Checks that there is no overlap between range and any other ChunkRange
+            // Specifically, checks that the greatest chunk less than or equal to range, if such a
+            // chunk exists, does not overlap with the min of range.
+            invariant(cr.getMax() <= range.getMin());
+        }
+    }
+
+    auto itHigh = _rangesToClean.lower_bound(range.getMin());
+    if (itHigh != _rangesToClean.end()) {
+        const ChunkRange& cr = itHigh->second;
+        // Checks that there is no overlap between range and any other ChunkRange
+        // Specifically, checks that the least chunk greater than or equal to range
+        // does not overlap with the max of range.
+        invariant(cr.getMin() >= range.getMax());
+    }
+
+    _rangesToClean.insert(std::make_pair(range.getMin(), range));
+}
+
+void MetadataManager::removeRangeToClean(const ChunkRange& range) {
+    auto it = _rangesToClean.upper_bound(range.getMin());
+    // We want our iterator to point at the greatest value
+    // that is still less than or equal to range.
+    if (it != _rangesToClean.begin()) {
+        --it;
+    }
+
+    for (; it != _rangesToClean.end() && it->second.getMin() < range.getMax();) {
+        if (it->second.getMax() <= range.getMin()) {
+            ++it;
+            continue;
+        }
+        // There's overlap between *it and range so we remove *it
+        // and then replace with new ranges.
+        ChunkRange oldChunk = it->second;
+        _rangesToClean.erase(it++);
+        if (oldChunk.getMin() < range.getMin()) {
+            ChunkRange newChunk = ChunkRange(oldChunk.getMin(), range.getMin());
+            addRangeToClean(newChunk);
+        }
+        if (oldChunk.getMax() > range.getMax()) {
+            ChunkRange newChunk = ChunkRange(range.getMax(), oldChunk.getMax());
+            addRangeToClean(newChunk);
+        }
+    }
+}
+
+void MetadataManager::append(BSONObjBuilder* builder) {
+    BSONArrayBuilder arr(builder->subarrayStart("rangesToClean"));
+    for (const auto& entry : _rangesToClean) {
+        BSONObjBuilder obj;
+        entry.second.append(&obj);
+        arr.append(obj.done());
+    }
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/s/metadata_manager.h b/src/mongo/db/s/metadata_manager.h
index d3625840acd..db0b3b1e255 100644
--- a/src/mongo/db/s/metadata_manager.h
+++ b/src/mongo/db/s/metadata_manager.h
@@ -25,6 +25,7 @@
  *    exception statement from all source files in the program, then also delete
  *    it in the license file.
  */
+
 #pragma once
 
 #include <list>
@@ -32,6 +33,7 @@
 
 #include "mongo/base/disallow_copying.h"
 #include "mongo/db/s/collection_metadata.h"
+#include "mongo/s/catalog/type_chunk.h"
 
 namespace mongo {
 
@@ -58,6 +60,26 @@ public:
      */
     void setActiveMetadata(std::unique_ptr<CollectionMetadata> newMetadata);
 
+    /**
+    * Adds a new range to be cleaned up.
+    * The newly introduced range must not overlap with the existing ranges.
+    */
+    void addRangeToClean(const ChunkRange& range);
+
+    /**
+     * Removes the specified range from the ranges to be cleaned up.
+     */
+    void removeRangeToClean(const ChunkRange& range);
+
+    /**
+     * Gets copy of _rangesToClean map (see below).
+     */
+    std::map<BSONObj, ChunkRange> getCopyOfRanges();
+
+    /*
+     * Appends information on all the chunk ranges in rangesToClean to builder.
+     */
+    void append(BSONObjBuilder* builder);
 
 private:
     friend class ScopedCollectionMetadata;
@@ -82,8 +104,12 @@ private:
     std::unique_ptr<CollectionMetadataTracker> _activeMetadataTracker;
 
     std::list<std::unique_ptr<CollectionMetadataTracker>> _metadataInUse;
-};
 
+    // Contains the information of which ranges of sharding keys need to
+    // be deleted from the shard. The map is from the minimum value of the
+    // range to be deleted (e.g. BSON("key" << 0)) to the entire chunk range.
+    std::map<BSONObj, ChunkRange> _rangesToClean;
+};
 
 class ScopedCollectionMetadata {
     MONGO_DISALLOW_COPYING(ScopedCollectionMetadata);
diff --git a/src/mongo/db/s/metadata_manager_test.cpp b/src/mongo/db/s/metadata_manager_test.cpp
index a1d872694c1..355c4c020d4 100644
--- a/src/mongo/db/s/metadata_manager_test.cpp
+++ b/src/mongo/db/s/metadata_manager_test.cpp
@@ -30,10 +30,13 @@
 
 #include "mongo/db/s/metadata_manager.h"
 
+#include "mongo/bson/bsonobjbuilder.h"
+#include "mongo/db/jsobj.h"
 #include "mongo/db/s/collection_metadata.h"
-#include "mongo/db/s/metadata_loader.h"
+#include "mongo/s/catalog/type_chunk.h"
 #include "mongo/stdx/memory.h"
 #include "mongo/unittest/unittest.h"
+#include "mongo/util/assert_util.h"
 
 namespace mongo {
 
@@ -68,5 +71,105 @@ TEST(MetadataManager, ResetActiveMetadata) {
     ASSERT_EQ(cm2Ptr, scopedMetadata2.getMetadata());
 };
 
+TEST(MetadataManager, AddAndRemoveRanges) {
+    MetadataManager mm;
+    ChunkRange cr1 = ChunkRange(BSON("key" << 0), BSON("key" << 10));
+    ChunkRange cr2 = ChunkRange(BSON("key" << 10), BSON("key" << 20));
+
+    mm.addRangeToClean(cr1);
+    ASSERT_EQ(mm.getCopyOfRanges().size(), 1UL);
+    mm.removeRangeToClean(cr1);
+    ASSERT_EQ(mm.getCopyOfRanges().size(), 0UL);
+
+    mm.addRangeToClean(cr1);
+    mm.addRangeToClean(cr2);
+    mm.removeRangeToClean(cr1);
+    ASSERT_EQ(mm.getCopyOfRanges().size(), 1UL);
+    auto ranges = mm.getCopyOfRanges();
+    ChunkRange remainingChunk = ranges.find(cr2.getMin())->second;
+    ASSERT_EQ(remainingChunk.toString(), cr2.toString());
+    mm.removeRangeToClean(cr2);
+}
+
+// Tests that a removal in the middle of an existing ChunkRange results in
+// two correct chunk ranges.
+TEST(MetadataManager, RemoveRangeInMiddleOfRange) {
+    MetadataManager mm;
+    ChunkRange cr1 = ChunkRange(BSON("key" << 0), BSON("key" << 10));
+
+    mm.addRangeToClean(cr1);
+    mm.removeRangeToClean(ChunkRange(BSON("key" << 4), BSON("key" << 6)));
+    ASSERT_EQ(mm.getCopyOfRanges().size(), 2UL);
+
+    auto ranges = mm.getCopyOfRanges();
+    auto it = ranges.find(BSON("key" << 0));
+    ChunkRange expectedChunk = ChunkRange(BSON("key" << 0), BSON("key" << 4));
+    ASSERT_EQ(it->second.toString(), expectedChunk.toString());
+
+    it++;
+    expectedChunk = ChunkRange(BSON("key" << 6), BSON("key" << 10));
+    ASSERT_EQ(it->second.toString(), expectedChunk.toString());
+
+    mm.removeRangeToClean(cr1);
+    ASSERT_EQ(mm.getCopyOfRanges().size(), 0UL);
+}
+
+// Tests removals that overlap with just one ChunkRange.
+TEST(MetadataManager, RemoveRangeWithSingleRangeOverlap) {
+    MetadataManager mm;
+    ChunkRange cr1 = ChunkRange(BSON("key" << 0), BSON("key" << 10));
+
+    mm.addRangeToClean(cr1);
+    mm.removeRangeToClean(ChunkRange(BSON("key" << 0), BSON("key" << 5)));
+    ASSERT_EQ(mm.getCopyOfRanges().size(), 1UL);
+    auto ranges = mm.getCopyOfRanges();
+    ChunkRange remainingChunk = ranges.find(BSON("key" << 5))->second;
+    ChunkRange expectedChunk = ChunkRange(BSON("key" << 5), BSON("key" << 10));
+    ASSERT_EQ(remainingChunk.toString(), expectedChunk.toString());
+
+    mm.removeRangeToClean(ChunkRange(BSON("key" << 4), BSON("key" << 6)));
+    ASSERT_EQ(mm.getCopyOfRanges().size(), 1UL);
+    ranges = mm.getCopyOfRanges();
+    remainingChunk = ranges.find(BSON("key" << 6))->second;
+    expectedChunk = ChunkRange(BSON("key" << 6), BSON("key" << 10));
+    ASSERT_EQ(remainingChunk.toString(), expectedChunk.toString());
+
+    mm.removeRangeToClean(ChunkRange(BSON("key" << 9), BSON("key" << 13)));
+    ASSERT_EQ(mm.getCopyOfRanges().size(), 1UL);
+    ranges = mm.getCopyOfRanges();
+    remainingChunk = ranges.find(BSON("key" << 6))->second;
+    expectedChunk = ChunkRange(BSON("key" << 6), BSON("key" << 9));
+    ASSERT_EQ(remainingChunk.toString(), expectedChunk.toString());
+
+    mm.removeRangeToClean(ChunkRange(BSON("key" << 0), BSON("key" << 10)));
+    ASSERT_EQ(mm.getCopyOfRanges().size(), 0UL);
+}
+
+// Tests removals that overlap with more than one ChunkRange.
+TEST(MetadataManager, RemoveRangeWithMultipleRangeOverlaps) {
+    MetadataManager mm;
+    ChunkRange cr1 = ChunkRange(BSON("key" << 0), BSON("key" << 10));
+    ChunkRange cr2 = ChunkRange(BSON("key" << 10), BSON("key" << 20));
+    ChunkRange cr3 = ChunkRange(BSON("key" << 20), BSON("key" << 30));
+
+    mm.addRangeToClean(cr1);
+    mm.addRangeToClean(cr2);
+    mm.addRangeToClean(cr3);
+    ASSERT_EQ(mm.getCopyOfRanges().size(), 3UL);
+
+    mm.removeRangeToClean(ChunkRange(BSON("key" << 8), BSON("key" << 22)));
+    ASSERT_EQ(mm.getCopyOfRanges().size(), 2UL);
+    auto ranges = mm.getCopyOfRanges();
+    auto it = ranges.find(BSON("key" << 0));
+    ChunkRange expectedChunk = ChunkRange(BSON("key" << 0), BSON("key" << 8));
+    ASSERT_EQ(it->second.toString(), expectedChunk.toString());
+    it++;
+    expectedChunk = ChunkRange(BSON("key" << 22), BSON("key" << 30));
+    ASSERT_EQ(it->second.toString(), expectedChunk.toString());
+
+    mm.removeRangeToClean(ChunkRange(BSON("key" << 0), BSON("key" << 30)));
+    ASSERT_EQ(mm.getCopyOfRanges().size(), 0UL);
+}
+
 }  // namespace
 }  // namespace mongo