SERVER-26370 Register incoming migrations on the ActiveMigrationsRegistry

In order to avoid race conditions where a shard could serve both as a
donor and recipient.

1 files changed, 49 insertions, 0 deletions

diff --git a/src/mongo/db/s/active_migrations_registry.h b/src/mongo/db/s/active_migrations_registry.h
index 06ffeaece07..e2f7ad06b8b 100644
--- a/src/mongo/db/s/active_migrations_registry.h
+++ b/src/mongo/db/s/active_migrations_registry.h
@@ -31,6 +31,7 @@
 #include <boost/optional.hpp>
 
 #include "mongo/base/disallow_copying.h"
+#include "mongo/db/s/migration_session_id.h"
 #include "mongo/s/move_chunk_request.h"
 #include "mongo/stdx/memory.h"
 #include "mongo/stdx/mutex.h"
@@ -40,6 +41,7 @@ namespace mongo {
 
 class OperationContext;
 class ScopedRegisterDonateChunk;
+class ScopedRegisterReceiveChunk;
 template <typename T>
 class StatusWith;
 
@@ -68,6 +70,15 @@ public:
     StatusWith<ScopedRegisterDonateChunk> registerDonateChunk(const MoveChunkRequest& args);
 
     /**
+     * If there are no migrations running on this shard, registers an active receive operation with
+     * the specified session id and returns a ScopedRegisterReceiveChunk, which will unregister it
+     * when it goes out of scope.
+     *
+     * Otherwise returns a ConflictingOperationInProgress error.
+     */
+    StatusWith<ScopedRegisterReceiveChunk> registerReceiveChunk(const NamespaceString& nss);
+
+    /**
      * If a migration has been previously registered through a call to registerDonateChunk returns
      * that namespace. Otherwise returns boost::none.
      */
@@ -83,6 +94,7 @@ public:
 
 private:
     friend class ScopedRegisterDonateChunk;
+    friend class ScopedRegisterReceiveChunk;
 
     // Describes the state of a currently active moveChunk operation
     struct ActiveMoveChunkState {
@@ -96,18 +108,36 @@ private:
         std::shared_ptr<Notification<Status>> notification;
     };
 
+    // Describes the state of a currently active receive chunk operation
+    struct ActiveReceiveChunkState {
+        ActiveReceiveChunkState(NamespaceString inNss) : nss(std::move(inNss)) {}
+
+        // Namesspace for which a chunk is being received
+        NamespaceString nss;
+    };
+
     /**
      * Unregisters a previously registered namespace with ongoing migration. Must only be called if
      * a previous call to registerDonateChunk has succeeded.
      */
     void _clearDonateChunk();
 
+    /**
+     * Unregisters a previously registered incoming migration. Must only be called if a previous
+     * call to registerReceiveChunk has succeeded.
+     */
+    void _clearReceiveChunk();
+
     // Protects the state below
     stdx::mutex _mutex;
 
     // If there is an active moveChunk operation going on, this field contains the request, which
     // initiated it
     boost::optional<ActiveMoveChunkState> _activeMoveChunkState;
+
+    // If there is an active receive of a chunk going on, this field contains the session id, which
+    // initiated it
+    boost::optional<ActiveReceiveChunkState> _activeReceiveChunkState;
 };
 
 /**
@@ -159,4 +189,23 @@ private:
     std::shared_ptr<Notification<Status>> _completionNotification;
 };
 
+/**
+ * Object of this class is returned from the registerReceiveChunk call of the active migrations
+ * registry.
+ */
+class ScopedRegisterReceiveChunk {
+    MONGO_DISALLOW_COPYING(ScopedRegisterReceiveChunk);
+
+public:
+    ScopedRegisterReceiveChunk(ActiveMigrationsRegistry* registry);
+    ~ScopedRegisterReceiveChunk();
+
+    ScopedRegisterReceiveChunk(ScopedRegisterReceiveChunk&&);
+    ScopedRegisterReceiveChunk& operator=(ScopedRegisterReceiveChunk&&);
+
+private:
+    // Registry from which to unregister the migration. Not owned.
+    ActiveMigrationsRegistry* _registry;
+};
+
 }  // namespace mongo