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.
   */
  #include "mongo/platform/basic.h"
  #include "mongo/platform/mutex.h"
  
  #include "mongo/executor/task_executor.h"
  
  namespace mongo {
  namespace executor {
  
  namespace {
  
  MONGO_FAIL_POINT_DEFINE(pauseScheduleCallWithCancelTokenUntilCanceled);
  
  /**
   * Provides exclusive access to an underlying Promise at set-time, guaranteeing that the Promise
   * will be set at most one time globally. This prevents races between completion and cancellation,
   * which normally would result in Promise throwing an invariant.
   */
  template <typename T>
  class ExclusivePromiseAccess {
  public:
      explicit ExclusivePromiseAccess(Promise<T>&& promise) : _promise(std::move(promise)) {}
  
      /**
       * Sets an error on the Promise if no result has been set; otherwise, does nothing.
       */
      template <typename Error>
      void setError(Error&& e) {
          if (!_completed.swap(true)) {
              _promise.setError(std::forward<Error>(e));
          }
      }
  
      /**
       * Sets a valid result on the Promise if no result has been set; otherwise, does nothing.
       */
      template <typename... Args>
      void emplaceValue(Args&&... args) {
          if (!_completed.swap(true)) {
              _promise.emplaceValue(std::forward<Args>(args)...);
          }
      }
  
  private:
      Promise<T> _promise;
      AtomicWord<bool> _completed;
  };
  
  template <typename T>
  Status wrapCallbackHandleWithCancelToken(
      const std::shared_ptr<TaskExecutor> executor,
      const StatusWith<TaskExecutor::CallbackHandle> swCallbackHandle,
      std::shared_ptr<ExclusivePromiseAccess<T>> promise,
      const CancellationToken& token) {
      if (!swCallbackHandle.isOK()) {
          return swCallbackHandle.getStatus();
      }
  
      token.onCancel()
          .unsafeToInlineFuture()
          .then(
              [executor, promise, callbackHandle = std::move(swCallbackHandle.getValue())]() mutable {
                  executor->cancel(callbackHandle);
                  promise->setError(TaskExecutor::kCallbackCanceledErrorStatus);
              })
          .getAsync([](auto) {});
      return Status::OK();
  }
  
  /**
   * Takes a schedule(Exhaust)RemoteCommand(OnAny)-style function and wraps it to return a future and
   * be cancelable with CancellationTokens.
   */
  template <typename Request, typename Response, typename ScheduleFn, typename CallbackFn>
  ExecutorFuture<Response> wrapScheduleCallWithCancelTokenAndFuture(
      const std::shared_ptr<TaskExecutor>& executor,
      ScheduleFn&& schedule,
      const Request& request,
      const CancellationToken& token,
      const BatonHandle& baton,
      const CallbackFn& cb) {
      if (token.isCanceled()) {
          return ExecutorFuture<Response>(executor, TaskExecutor::kCallbackCanceledErrorStatus);
      }
  
      auto [promise, future] = makePromiseFuture<Response>();
  
      // This has to be made shared because otherwise we'd have to move the access into this
      // callback, and would be unable to use it in the case where scheduling the request fails below.
      auto exclusivePromiseAccess =
          std::make_shared<ExclusivePromiseAccess<Response>>(std::move(promise));
      auto signalPromiseOnCompletion = [exclusivePromiseAccess,
                                        cb = std::move(cb)](const auto& args) mutable {
          // Upon completion, unconditionally run our callback.
          cb(args);
          auto status = args.response.status;
  
          // Only mark the future as complete when the moreToCome flag is false, even if an error has
          // occured.
          if (!args.response.moreToCome) {
              if (status.isOK()) {
                  exclusivePromiseAccess->emplaceValue(std::move(args.response));
              } else {
                  // Only set an error on failures to send the request (including if the request was
                  // canceled). Errors from the remote host will be contained in the response.
                  exclusivePromiseAccess->setError(status);
              }
          }
      };
  
      // Fail point to make this method to wait until the token is canceled.
      if (!token.isCanceled()) {
          try {
              pauseScheduleCallWithCancelTokenUntilCanceled.pauseWhileSetAndNotCanceled(
                  Interruptible::notInterruptible(), token);
          } catch (ExceptionFor<ErrorCodes::Interrupted>&) {
              // Swallow the interrupted exception that arrives from canceling a failpoint.
          }
      }
  
      auto scheduleStatus = wrapCallbackHandleWithCancelToken(
          executor,
          std::forward<ScheduleFn>(schedule)(request, std::move(signalPromiseOnCompletion), baton),
          exclusivePromiseAccess,
          token);
  
      if (!scheduleStatus.isOK()) {
          // If scheduleStatus is not okay, then the callback signalPromiseOnCompletion should never
          // run, meaning that it will be okay to set the promise here.
          exclusivePromiseAccess->setError(scheduleStatus);
      }
  
      return std::move(future).thenRunOn(executor);
  }
  }  // namespace
  
  TaskExecutor::TaskExecutor() = default;
  TaskExecutor::~TaskExecutor() = default;
  
  void TaskExecutor::schedule(OutOfLineExecutor::Task func) {
      auto cb = CallbackFn([func = std::move(func)](const CallbackArgs& args) { func(args.status); });
      auto statusWithCallback = scheduleWork(std::move(cb));
      if (!statusWithCallback.isOK()) {
          // The callback was not scheduled or moved from, it is still valid. Run it inline to inform
          // it of the error. Construct a CallbackArgs for it, only CallbackArgs::status matters here.
          CallbackArgs args(this, {}, statusWithCallback.getStatus(), nullptr);
          invariant(cb);  // NOLINT(bugprone-use-after-move)
          cb(args);       // NOLINT(bugprone-use-after-move)
      }
  }
  
  ExecutorFuture<void> TaskExecutor::sleepUntil(Date_t when, const CancellationToken& token) {
      if (token.isCanceled()) {
          return ExecutorFuture<void>(shared_from_this(), TaskExecutor::kCallbackCanceledErrorStatus);
      }
  
      if (when <= now()) {
          return ExecutorFuture<void>(shared_from_this());
      }
  
      /**
       * Encapsulates the promise associated with the result future.
       */
      struct AlarmState {
          void signal(const Status& status) {
              if (status.isOK()) {
                  promise->emplaceValue();
              } else {
                  promise->setError(status);
              }
          }
  
          std::shared_ptr<ExclusivePromiseAccess<void>> promise;
      };
  
      auto [promise, future] = makePromiseFuture<void>();
      // This has to be shared because Promises (and therefore AlarmState) are move-only and we need
      // to maintain two copies: One to capture in the scheduleWorkAt callback, and one locally in
      // case scheduling the request fails.
      auto exclusivePromiseAccess =
          std::make_shared<ExclusivePromiseAccess<void>>(std::move(promise));
      auto alarmState = std::make_shared<AlarmState>(AlarmState{exclusivePromiseAccess});
  
      // Schedule a task to signal the alarm when the deadline is reached.
      auto cbHandle = scheduleWorkAt(
          when, [alarmState](const auto& args) mutable { alarmState->signal(args.status); });
  
      // Handle cancellation via the input CancellationToken.
      auto scheduleStatus = wrapCallbackHandleWithCancelToken(
          shared_from_this(), std::move(cbHandle), exclusivePromiseAccess, token);
  
      if (!scheduleStatus.isOK()) {
          // If scheduleStatus is not okay, then the callback passed to scheduleWorkAt should never
          // run, meaning that it will be okay to set the promise here.
          alarmState->signal(scheduleStatus);
      }
  
      // TODO (SERVER-51285): Optimize to avoid an additional call to schedule to run the callback
      // chained by the caller of sleepUntil.
      return std::move(future).thenRunOn(shared_from_this());
  }
  
  
  TaskExecutor::CallbackState::CallbackState() = default;
  TaskExecutor::CallbackState::~CallbackState() = default;
  
  TaskExecutor::CallbackHandle::CallbackHandle() = default;
  TaskExecutor::CallbackHandle::CallbackHandle(std::shared_ptr<CallbackState> callback)
      : _callback(std::move(callback)) {}
  
  TaskExecutor::EventState::EventState() = default;
  TaskExecutor::EventState::~EventState() = default;
  
  TaskExecutor::EventHandle::EventHandle() = default;
  TaskExecutor::EventHandle::EventHandle(std::shared_ptr<EventState> event)
      : _event(std::move(event)) {}
  
  TaskExecutor::CallbackArgs::CallbackArgs(TaskExecutor* theExecutor,
                                           CallbackHandle theHandle,
                                           Status theStatus,
                                           OperationContext* theTxn)
      : executor(theExecutor),
        myHandle(std::move(theHandle)),
        status(std::move(theStatus)),
        opCtx(theTxn) {}
  
  
  TaskExecutor::RemoteCommandCallbackArgs::RemoteCommandCallbackArgs(
      TaskExecutor* theExecutor,
      const CallbackHandle& theHandle,
      const RemoteCommandRequest& theRequest,
      const ResponseStatus& theResponse)
      : executor(theExecutor), myHandle(theHandle), request(theRequest), response(theResponse) {}
  
  TaskExecutor::RemoteCommandCallbackArgs::RemoteCommandCallbackArgs(
      const RemoteCommandOnAnyCallbackArgs& other, size_t idx)
      : executor(other.executor),
        myHandle(other.myHandle),
        request(other.request, idx),
        response(other.response) {}
  
  TaskExecutor::RemoteCommandOnAnyCallbackArgs::RemoteCommandOnAnyCallbackArgs(
      TaskExecutor* theExecutor,
      const CallbackHandle& theHandle,
      const RemoteCommandRequestOnAny& theRequest,
      const ResponseOnAnyStatus& theResponse)
      : executor(theExecutor), myHandle(theHandle), request(theRequest), response(theResponse) {}
  
  TaskExecutor::CallbackState* TaskExecutor::getCallbackFromHandle(const CallbackHandle& cbHandle) {
      return cbHandle.getCallback();
  }
  
  TaskExecutor::EventState* TaskExecutor::getEventFromHandle(const EventHandle& eventHandle) {
      return eventHandle.getEvent();
  }
  
  void TaskExecutor::setEventForHandle(EventHandle* eventHandle, std::shared_ptr<EventState> event) {
      eventHandle->setEvent(std::move(event));
  }
  
  void TaskExecutor::setCallbackForHandle(CallbackHandle* cbHandle,
                                          std::shared_ptr<CallbackState> callback) {
      cbHandle->setCallback(std::move(callback));
  }
  
  
  StatusWith<TaskExecutor::CallbackHandle> TaskExecutor::scheduleRemoteCommand(
      const RemoteCommandRequest& request,
      const RemoteCommandCallbackFn& cb,
      const BatonHandle& baton) {
      return scheduleRemoteCommandOnAny(request,
                                        [cb](const RemoteCommandOnAnyCallbackArgs& args) {
                                            cb({args, 0});
                                        },
                                        baton);
  }
  
  ExecutorFuture<TaskExecutor::ResponseStatus> TaskExecutor::scheduleRemoteCommand(
      const RemoteCommandRequest& request, const CancellationToken& token, const BatonHandle& baton) {
      return wrapScheduleCallWithCancelTokenAndFuture<decltype(request),
                                                      TaskExecutor::ResponseStatus>(
          shared_from_this(),
          [this](const auto&... args) { return scheduleRemoteCommand(args...); },
          request,
          token,
          baton,
          [](const auto& args) {});
  }
  
  ExecutorFuture<TaskExecutor::ResponseOnAnyStatus> TaskExecutor::scheduleRemoteCommandOnAny(
      const RemoteCommandRequestOnAny& request,
      const CancellationToken& token,
      const BatonHandle& baton) {
      return wrapScheduleCallWithCancelTokenAndFuture<decltype(request),
                                                      TaskExecutor::ResponseOnAnyStatus>(
          shared_from_this(),
          [this](const auto&... args) { return scheduleRemoteCommandOnAny(args...); },
          request,
          token,
          baton,
          [](const auto& args) {});
  }
  
  StatusWith<TaskExecutor::CallbackHandle> TaskExecutor::scheduleExhaustRemoteCommand(
      const RemoteCommandRequest& request,
      const RemoteCommandCallbackFn& cb,
      const BatonHandle& baton) {
      return scheduleExhaustRemoteCommandOnAny(request,
                                               [cb](const RemoteCommandOnAnyCallbackArgs& args) {
                                                   cb({args, 0});
                                               },
                                               baton);
  }
  
  ExecutorFuture<TaskExecutor::ResponseStatus> TaskExecutor::scheduleExhaustRemoteCommand(
      const RemoteCommandRequest& request,
      const RemoteCommandCallbackFn& cb,
      const CancellationToken& token,
      const BatonHandle& baton) {
      return wrapScheduleCallWithCancelTokenAndFuture<decltype(request),
                                                      TaskExecutor::ResponseStatus>(
          shared_from_this(),
          [this](const auto&... args) { return scheduleExhaustRemoteCommand(args...); },
          request,
          token,
          baton,
          cb);
  }
  
  ExecutorFuture<TaskExecutor::ResponseOnAnyStatus> TaskExecutor::scheduleExhaustRemoteCommandOnAny(
      const RemoteCommandRequestOnAny& request,
      const RemoteCommandOnAnyCallbackFn& cb,
      const CancellationToken& token,
      const BatonHandle& baton) {
      return wrapScheduleCallWithCancelTokenAndFuture<decltype(request),
                                                      TaskExecutor::ResponseOnAnyStatus>(
          shared_from_this(),
          [this](const auto&... args) { return scheduleExhaustRemoteCommandOnAny(args...); },
          request,
          token,
          baton,
          cb);
  }
  }  // namespace executor
  }  // namespace mongo
  