SERVER-24629 Add concurrency protection to MetadataManager

2 files changed, 33 insertions, 5 deletions

diff --git a/src/mongo/db/s/metadata_manager.cpp b/src/mongo/db/s/metadata_manager.cpp
index 9547bd64822..cf2cb76def5 100644
--- a/src/mongo/db/s/metadata_manager.cpp
+++ b/src/mongo/db/s/metadata_manager.cpp
@@ -36,19 +36,26 @@ namespace mongo {
 
 MetadataManager::MetadataManager() = default;
 
+MetadataManager::~MetadataManager() {
+    stdx::lock_guard<stdx::mutex> scopedLock(_managerLock);
+    invariant(!_activeMetadataTracker || _activeMetadataTracker->usageCounter == 0);
+}
+
 ScopedCollectionMetadata MetadataManager::getActiveMetadata() {
+    stdx::lock_guard<stdx::mutex> scopedLock(_managerLock);
     invariant(_activeMetadataTracker);
     return ScopedCollectionMetadata(this, _activeMetadataTracker.get());
 }
 
 void MetadataManager::setActiveMetadata(std::unique_ptr<CollectionMetadata> newMetadata) {
+    stdx::lock_guard<stdx::mutex> scopedLock(_managerLock);
     if (_activeMetadataTracker && _activeMetadataTracker->usageCounter > 0) {
         _metadataInUse.push_front(std::move(_activeMetadataTracker));
     }
     _activeMetadataTracker = stdx::make_unique<CollectionMetadataTracker>(std::move(newMetadata));
 }
 
-void MetadataManager::_removeMetadata(CollectionMetadataTracker* metadataTracker) {
+void MetadataManager::_removeMetadata_inlock(CollectionMetadataTracker* metadataTracker) {
     invariant(metadataTracker->usageCounter == 0);
     auto i = _metadataInUse.begin();
     auto e = _metadataInUse.end();
@@ -65,6 +72,7 @@ MetadataManager::CollectionMetadataTracker::CollectionMetadataTracker(
     std::unique_ptr<CollectionMetadata> m)
     : metadata(std::move(m)), usageCounter(0){};
 
+// called in lock
 ScopedCollectionMetadata::ScopedCollectionMetadata(
     MetadataManager* manager, MetadataManager::CollectionMetadataTracker* tracker)
     : _manager(manager), _tracker(tracker) {
@@ -72,9 +80,10 @@ ScopedCollectionMetadata::ScopedCollectionMetadata(
 }
 
 ScopedCollectionMetadata::~ScopedCollectionMetadata() {
+    stdx::lock_guard<stdx::mutex> scopedLock(_manager->_managerLock);
     invariant(_tracker->usageCounter > 0);
     if (--_tracker->usageCounter == 0) {
-        _manager->_removeMetadata(_tracker);
+        _manager->_removeMetadata_inlock(_tracker);
     }
 }
 
@@ -102,10 +111,16 @@ ScopedCollectionMetadata& ScopedCollectionMetadata::operator=(ScopedCollectionMe
 }
 
 std::map<BSONObj, ChunkRange> MetadataManager::getCopyOfRanges() {
+    stdx::lock_guard<stdx::mutex> scopedLock(_managerLock);
     return _rangesToClean;
 }
 
 void MetadataManager::addRangeToClean(const ChunkRange& range) {
+    stdx::lock_guard<stdx::mutex> scopedLock(_managerLock);
+    _addRangeToClean_inlock(range);
+}
+
+void MetadataManager::_addRangeToClean_inlock(const ChunkRange& range) {
     auto itLow = _rangesToClean.upper_bound(range.getMin());
     if (itLow != _rangesToClean.begin()) {
         --itLow;
@@ -134,6 +149,11 @@ void MetadataManager::addRangeToClean(const ChunkRange& range) {
 }
 
 void MetadataManager::removeRangeToClean(const ChunkRange& range) {
+    stdx::lock_guard<stdx::mutex> scopedLock(_managerLock);
+    _removeRangeToClean_inlock(range);
+}
+
+void MetadataManager::_removeRangeToClean_inlock(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.
@@ -152,17 +172,18 @@ void MetadataManager::removeRangeToClean(const ChunkRange& range) {
         _rangesToClean.erase(it++);
         if (oldChunk.getMin() < range.getMin()) {
             ChunkRange newChunk = ChunkRange(oldChunk.getMin(), range.getMin());
-            addRangeToClean(newChunk);
+            _addRangeToClean_inlock(newChunk);
         }
         if (oldChunk.getMax() > range.getMax()) {
             ChunkRange newChunk = ChunkRange(range.getMax(), oldChunk.getMax());
-            addRangeToClean(newChunk);
+            _addRangeToClean_inlock(newChunk);
         }
     }
 }
 
 void MetadataManager::append(BSONObjBuilder* builder) {
     BSONArrayBuilder arr(builder->subarrayStart("rangesToClean"));
+    stdx::lock_guard<stdx::mutex> scopedLock(_managerLock);
     for (const auto& entry : _rangesToClean) {
         BSONObjBuilder obj;
         entry.second.append(&obj);
diff --git a/src/mongo/db/s/metadata_manager.h b/src/mongo/db/s/metadata_manager.h
index db0b3b1e255..73f91f7e44f 100644
--- a/src/mongo/db/s/metadata_manager.h
+++ b/src/mongo/db/s/metadata_manager.h
@@ -34,6 +34,7 @@
 #include "mongo/base/disallow_copying.h"
 #include "mongo/db/s/collection_metadata.h"
 #include "mongo/s/catalog/type_chunk.h"
+#include "mongo/stdx/memory.h"
 
 namespace mongo {
 
@@ -44,6 +45,7 @@ class MetadataManager {
 
 public:
     MetadataManager();
+    ~MetadataManager();
 
     /**
      * An ActiveMetadata must be set before this function can be called.
@@ -99,7 +101,10 @@ private:
      * Removes the CollectionMetadata stored in the tracker from the _metadataInUse
      * list (if it's there).
      */
-    void _removeMetadata(CollectionMetadataTracker* metadataTracker);
+    void _removeMetadata_inlock(CollectionMetadataTracker* metadataTracker);
+
+    void _addRangeToClean_inlock(const ChunkRange& range);
+    void _removeRangeToClean_inlock(const ChunkRange& range);
 
     std::unique_ptr<CollectionMetadataTracker> _activeMetadataTracker;
 
@@ -109,6 +114,8 @@ private:
     // 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;
+
+    stdx::mutex _managerLock;
 };
 
 class ScopedCollectionMetadata {