text: /**
   *    Copyright (C) 2018-present MongoDB, Inc.
   *
   *    This program is free software: you can redistribute it and/or modify
   *    it under the terms of the Server Side Public License, version 1,
   *    as published by MongoDB, Inc.
   *
   *    This program is distributed in the hope that it will be useful,
   *    but WITHOUT ANY WARRANTY; without even the implied warranty of
   *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   *    Server Side Public License for more details.
   *
   *    You should have received a copy of the Server Side Public License
   *    along with this program. If not, see
   *    <http://www.mongodb.com/licensing/server-side-public-license>.
   *
   *    As a special exception, the copyright holders give permission to link the
   *    code of portions of this program with the OpenSSL library under certain
   *    conditions as described in each individual source file and distribute
   *    linked combinations including the program with the OpenSSL library. You
   *    must comply with the Server Side Public License in all respects for
   *    all of the code used other than as permitted herein. If you modify file(s)
   *    with this exception, you may extend this exception to your version of the
   *    file(s), but you are not obligated to do so. If you do not wish to do so,
   *    delete this exception statement from your version. If you delete this
   *    exception statement from all source files in the program, then also delete
   *    it in the license file.
   */
  
  #pragma once
  
  #include <list>
  #include <map>
  #include <vector>
  
  #include "mongo/bson/bsonobj.h"
  #include "mongo/db/namespace_string.h"
  #include "mongo/db/s/balancer/balancer_policy.h"
  #include "mongo/db/s/balancer/type_migration.h"
  #include "mongo/executor/task_executor.h"
  #include "mongo/platform/mutex.h"
  #include "mongo/s/catalog/dist_lock_manager.h"
  #include "mongo/s/request_types/migration_secondary_throttle_options.h"
  #include "mongo/stdx/condition_variable.h"
  #include "mongo/stdx/unordered_map.h"
  #include "mongo/util/concurrency/notification.h"
  #include "mongo/util/concurrency/with_lock.h"
  
  namespace mongo {
  
  class OperationContext;
  class ScopedMigrationRequest;
  class ServiceContext;
  class Status;
  template <typename T>
  class StatusWith;
  
  // Uniquely identifies a migration, regardless of shard and version.
  typedef std::string MigrationIdentifier;
  typedef std::map<MigrationIdentifier, Status> MigrationStatuses;
  
  /**
   * Manages and executes parallel migrations for the balancer.
   */
  class MigrationManager {
      MigrationManager(const MigrationManager&) = delete;
      MigrationManager& operator=(const MigrationManager&) = delete;
  
  public:
      MigrationManager(ServiceContext* serviceContext);
      ~MigrationManager();
  
      /**
       * A blocking method that attempts to schedule all the migrations specified in
       * "candidateMigrations" and wait for them to complete. Takes the distributed lock for each
       * collection with a chunk being migrated.
       *
       * If any of the migrations, which were scheduled in parallel fails with a LockBusy error
       * reported from the shard, retries it serially without the distributed lock.
       *
       * Returns a map of migration Status objects to indicate the success/failure of each migration.
       */
      MigrationStatuses executeMigrationsForAutoBalance(
          OperationContext* opCtx,
          const std::vector<MigrateInfo>& migrateInfos,
          uint64_t maxChunkSizeBytes,
          const MigrationSecondaryThrottleOptions& secondaryThrottle,
          bool waitForDelete);
  
      /**
       * A blocking method that attempts to schedule the migration specified in "migrateInfo" and
       * waits for it to complete. Takes the distributed lock for the namespace which is being
       * migrated.
       *
       * Returns the status of the migration.
       */
      Status executeManualMigration(OperationContext* opCtx,
                                    const MigrateInfo& migrateInfo,
                                    uint64_t maxChunkSizeBytes,
                                    const MigrationSecondaryThrottleOptions& secondaryThrottle,
                                    bool waitForDelete);
  
      /**
       * Non-blocking method that puts the migration manager in the kRecovering state, in which
       * new migration requests will block until finishRecovery is called. Then reacquires distributed
       * locks for the balancer and any active migrations. The distributed locks are taken with local
       * write concern, since this is called in drain mode where majority writes are not yet possible.
       *
       * The active migration recovery may fail and be abandoned, setting the state to kEnabled.
       */
      void startRecoveryAndAcquireDistLocks(OperationContext* opCtx);
  
      /**
       * Blocking method that must only be called after startRecovery has been called. Recovers the
       * state of the migration manager (if necessary and able) and puts it in the kEnabled state,
       * where it will accept new migrations. Any migrations waiting on the recovery state will be
       * unblocked once the state is kEnabled, and then this function waits for the recovered active
       * migrations to finish before returning.
       *
       * The active migration recovery may fail and be abandoned, setting the state to kEnabled and
       * unblocking any process waiting on the recovery state.
       */
      void finishRecovery(OperationContext* opCtx,
                          uint64_t maxChunkSizeBytes,
                          const MigrationSecondaryThrottleOptions& secondaryThrottle);
  
      /**
       * Non-blocking method that should never be called concurrently with finishRecovery. Puts the
       * manager in a state where all subsequently scheduled migrations will immediately fail (without
       * ever getting scheduled) and all active ones will be cancelled. It has no effect if the
       * migration manager is already stopping or stopped.
       */
      void interruptAndDisableMigrations();
  
      /**
       * Blocking method that waits for any currently scheduled migrations to complete. Must be
       * called after interruptAndDisableMigrations has been called in order to be able to re-enable
       * migrations again.
       */
      void drainActiveMigrations();
  
  private:
      // The current state of the migration manager
      enum class State {  // Allowed transitions:
          kStopped,       // kRecovering
          kRecovering,    // kEnabled, kStopping
          kEnabled,       // kStopping
          kStopping,      // kStopped
      };
  
      /**
       * Tracks the execution state of a single migration.
       */
      struct Migration {
          Migration(NamespaceString nss, BSONObj moveChunkCmdObj);
          ~Migration();
  
          // Namespace for which this migration applies
          NamespaceString nss;
  
          // Command object representing the migration
          BSONObj moveChunkCmdObj;
  
          // Callback handle for the migration network request. If the migration has not yet been sent
          // on the network, this value is not set.
          boost::optional<executor::TaskExecutor::CallbackHandle> callbackHandle;
  
          // Notification, which will be signaled when the migration completes
          std::shared_ptr<Notification<executor::RemoteCommandResponse>> completionNotification;
      };
  
      // Used as a type in which to store a list of active migrations. The reason to choose list is
      // that its iterators do not get invalidated when entries are removed around them. This allows
      // O(1) removal time.
      using MigrationsList = std::list<Migration>;
  
      using CollectionMigrationsStateMap = stdx::unordered_map<NamespaceString, MigrationsList>;
  
      using ScopedMigrationRequestsMap =
          std::map<MigrationIdentifier, StatusWith<ScopedMigrationRequest>>;
  
      /**
       * Optionally takes the collection distributed lock and schedules a chunk migration with the
       * specified parameters. May block for distributed lock acquisition. If dist lock acquisition is
       * successful (or not done), schedules the migration request and returns a notification which
       * can be used to obtain the outcome of the operation.
       */
      std::shared_ptr<Notification<executor::RemoteCommandResponse>> _schedule(
          OperationContext* opCtx,
          const MigrateInfo& migrateInfo,
          uint64_t maxChunkSizeBytes,
          const MigrationSecondaryThrottleOptions& secondaryThrottle,
          bool waitForDelete,
          ScopedMigrationRequestsMap* scopedMigrationRequests);
  
      /**
       * Acquires the collection distributed lock for the specified namespace and if it succeeds,
       * schedules the migration.
       *
       * The distributed lock is acquired before scheduling the first migration for the collection and
       * is only released when all active migrations on the collection have finished.
       *
       * Assumes that the migration document has already been written if no ScopedMigrationRequestsMap
       * pointer is passed. Otherwise, writes the migration document under the collection distributed
       * lock and adds it to the map.
       */
      void _schedule(WithLock,
                     OperationContext* opCtx,
                     const HostAndPort& targetHost,
                     Migration migration,
                     const MigrateInfo& migrateInfo,
                     bool waitForDelete,
                     ScopedMigrationRequestsMap* scopedMigrationRequests);
  
      /**
       * Used internally for migrations scheduled with the distributed lock acquired by the config
       * server. Called exactly once for each scheduled migration, it will signal the migration in the
       * passed iterator and if this is the last migration for the collection will free the collection
       * distributed lock.
       */
      void _complete(WithLock,
                     OperationContext* opCtx,
                     MigrationsList::iterator itMigration,
                     const executor::RemoteCommandResponse& remoteCommandResponse);
  
      /**
       * If the state of the migration manager is kStopping, checks whether there are any outstanding
       * scheduled requests and if there aren't any signals the class condition variable.
       */
      void _checkDrained(WithLock);
  
      /**
       * Blocking call, which waits for the migration manager to leave the recovering state (if it is
       * currently recovering).
       */
      void _waitForRecovery();
  
      /**
       * Should only be called from startRecovery or finishRecovery functions when the migration
       * manager is in either the kStopped or kRecovering state. Releases all the distributed locks
       * that the balancer holds, clears the config.migrations collection, changes the state of the
       * migration manager to kEnabled. Then unblocks all processes waiting for kEnabled state.
       */
      void _abandonActiveMigrationsAndEnableManager(OperationContext* opCtx);
  
      /**
       * Parses a moveChunk RemoteCommandResponse's two levels of Status objects and distiguishes
       * between errors generated by this config server and the shard primary to which the moveChunk
       * command was sent.
       *
       * If the command failed because of stepdown of this config server, the migration document
       * managed by 'scopedMigrationRequest' is saved for later balancer recovery and a
       * BalancerInterrupted error is returned. If the command failed because the shard to which the
       * command was sent returned an error, the migration document is not saved and the error is
       * returned without conversion.
       */
      Status _processRemoteCommandResponse(
          const executor::RemoteCommandResponse& remoteCommandResponse,
          ScopedMigrationRequest* scopedMigrationRequest);
  
      // The service context under which this migration manager runs.
      ServiceContext* const _serviceContext;
  
      // Used as a constant session ID for all distributed locks that this MigrationManager holds.
      // Currently required so that locks can be reacquired for the balancer in startRecovery and then
      // overtaken in later operations.
      const OID _lockSessionID{OID::gen()};
  
      // Carries migration information over from startRecovery to finishRecovery. Should only be set
      // in startRecovery and then accessed in finishRecovery.
      stdx::unordered_map<NamespaceString, std::list<MigrationType>> _migrationRecoveryMap;
  
      // Protects the class state below.
      Mutex _mutex = MONGO_MAKE_LATCH("MigrationManager::_mutex");
  
      // Always start the migration manager in a stopped state.
      State _state{State::kStopped};
  
      // Condition variable, which is waited on when the migration manager's state is changing and
      // signaled when the state change is complete.
      stdx::condition_variable _condVar;
  
      // Maps collection namespaces to that collection's active migrations.
      CollectionMigrationsStateMap _activeMigrations;
  };
  
  }  // namespace mongo
  