Core.h

Go to the documentation of this file.

/*
 * Copyright 2014-present Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <atomic>
#include <mutex>
#include <stdexcept>
#include <utility>
#include <vector>

#include <folly/Executor.h>
#include <folly/Function.h>
#include <folly/Optional.h>
#include <folly/ScopeGuard.h>
#include <folly/Try.h>
#include <folly/Utility.h>
#include <folly/lang/Assume.h>
#include <folly/lang/Exception.h>
#include <folly/synchronization/MicroSpinLock.h>
#include <glog/logging.h>

#include <folly/io/async/Request.h>

namespace folly {
namespace futures {
namespace detail {

enum class State : uint8_t {
  Start = 1 << 0,
  OnlyResult = 1 << 1,
  OnlyCallback = 1 << 2,
  Proxy = 1 << 3,
  Done = 1 << 4,
  Empty = 1 << 5,
};
constexpr State operator&(State a, State b) {
  return State(uint8_t(a) & uint8_t(b));
}
constexpr State operator|(State a, State b) {
  return State(uint8_t(a) | uint8_t(b));
}
constexpr State operator^(State a, State b) {
  return State(uint8_t(a) ^ uint8_t(b));
}
constexpr State operator~(State a) {
  return State(~uint8_t(a));
}

struct SpinLock : private MicroSpinLock {
  SpinLock() : MicroSpinLock{0} {}

  using MicroSpinLock::lock;
  using MicroSpinLock::unlock;
};
static_assert(sizeof(SpinLock) == 1, "missized");

template <typename T>
class Core final {
  static_assert(
      !std::is_void<T>::value,
      "void futures are not supported. Use Unit instead.");

 public:
  using Result = Try<T>;
  using Callback = folly::Function<void(Result&&)>;

  static Core* make() {
    return new Core();
  }

  static Core* make(Try<T>&& t) {
    return new Core(std::move(t));
  }

  template <typename... Args>
  static Core<T>* make(in_place_t, Args&&... args) {
    return new Core<T>(in_place, std::forward<Args>(args)...);
  }

  // not copyable
  Core(Core const&) = delete;
  Core& operator=(Core const&) = delete;

  // not movable (see comment in the implementation of Future::then)
  Core(Core&&) noexcept = delete;
  Core& operator=(Core&&) = delete;

  bool hasCallback() const noexcept {
    constexpr auto allowed = State::OnlyCallback | State::Done;
    auto const state = state_.load(std::memory_order_acquire);
    return State() != (state & allowed);
  }

  bool hasResult() const noexcept {
    constexpr auto allowed = State::OnlyResult | State::Done;
    auto core = this;
    auto state = core->state_.load(std::memory_order_acquire);
    while (state == State::Proxy) {
      core = core->proxy_;
      state = core->state_.load(std::memory_order_acquire);
    }
    return State() != (state & allowed);
  }

  bool ready() const noexcept {
    return hasResult();
  }

  Try<T>& getTry() {
    DCHECK(hasResult());
    auto core = this;
    while (core->state_.load(std::memory_order_relaxed) == State::Proxy) {
      core = core->proxy_;
    }
    return core->result_;
  }
  Try<T> const& getTry() const {
    DCHECK(hasResult());
    auto core = this;
    while (core->state_.load(std::memory_order_relaxed) == State::Proxy) {
      core = core->proxy_;
    }
    return core->result_;
  }

  template <typename F>
  void setCallback(F&& func, std::shared_ptr<folly::RequestContext> context) {
    DCHECK(!hasCallback());

    // construct callback_ first; if that fails, context_ will not leak
    ::new (&callback_) Callback(std::forward<F>(func));
    ::new (&context_) Context(std::move(context));

    auto state = state_.load(std::memory_order_acquire);

    if (state == State::Start) {
      if (state_.compare_exchange_strong(
              state, State::OnlyCallback, std::memory_order_release)) {
        return;
      }
      assume(state == State::OnlyResult || state == State::Proxy);
    }

    if (state == State::OnlyResult) {
      state_.store(State::Done, std::memory_order_relaxed);
      doCallback();
      return;
    }

    if (state == State::Proxy) {
      return proxyCallback();
    }

    terminate_with<std::logic_error>("setCallback unexpected state");
  }

  void setProxy(Core* proxy) {
    DCHECK(!hasResult());

    proxy_ = proxy;

    auto state = state_.load(std::memory_order_acquire);
    switch (state) {
      case State::Start:
        if (state_.compare_exchange_strong(
                state, State::Proxy, std::memory_order_release)) {
          break;
        }
        assume(state == State::OnlyCallback);
        FOLLY_FALLTHROUGH;

      case State::OnlyCallback:
        proxyCallback();
        break;

      default:
        terminate_with<std::logic_error>("setCallback unexpected state");
    }

    detachOne();
  }

  void setResult(Try<T>&& t) {
    DCHECK(!hasResult());

    ::new (&result_) Result(std::move(t));

    auto state = state_.load(std::memory_order_acquire);
    while (true) {
      switch (state) {
        case State::Start:
          if (state_.compare_exchange_strong(
                  state, State::OnlyResult, std::memory_order_release)) {
            return;
          }
          assume(state == State::OnlyCallback);
          FOLLY_FALLTHROUGH;

        case State::OnlyCallback:
          if (state_.compare_exchange_strong(
                  state, State::Done, std::memory_order_release)) {
            doCallback();
            return;
          }
          FOLLY_FALLTHROUGH;

        default:
          terminate_with<std::logic_error>("setResult unexpected state");
      }
    }
  }

  void detachFuture() noexcept {
    detachOne();
  }

  void detachPromise() noexcept {
    DCHECK(hasResult());
    detachOne();
  }

  void setExecutor(
      Executor::KeepAlive<> x,
      int8_t priority = Executor::MID_PRI) {
    DCHECK(state_ != State::OnlyCallback);
    executor_ = std::move(x);
    priority_ = priority;
  }

  void setExecutor(Executor* x, int8_t priority = Executor::MID_PRI) {
    setExecutor(getKeepAliveToken(x), priority);
  }

  Executor* getExecutor() const {
    return executor_.get();
  }

  int8_t getPriority() const {
    return priority_;
  }

  void raise(exception_wrapper e) {
    std::lock_guard<SpinLock> lock(interruptLock_);
    if (!interrupt_ && !hasResult()) {
      interrupt_ = std::make_unique<exception_wrapper>(std::move(e));
      if (interruptHandler_) {
        interruptHandler_(*interrupt_);
      }
    }
  }

  std::function<void(exception_wrapper const&)> getInterruptHandler() {
    if (!interruptHandlerSet_.load(std::memory_order_acquire)) {
      return nullptr;
    }
    std::lock_guard<SpinLock> lock(interruptLock_);
    return interruptHandler_;
  }

  template <typename F>
  void setInterruptHandler(F&& fn) {
    std::lock_guard<SpinLock> lock(interruptLock_);
    if (!hasResult()) {
      if (interrupt_) {
        fn(as_const(*interrupt_));
      } else {
        setInterruptHandlerNoLock(std::forward<F>(fn));
      }
    }
  }

  void setInterruptHandlerNoLock(
      std::function<void(exception_wrapper const&)> fn) {
    interruptHandlerSet_.store(true, std::memory_order_relaxed);
    interruptHandler_ = std::move(fn);
  }

 private:
  Core() : state_(State::Start), attached_(2) {}

  explicit Core(Try<T>&& t)
      : result_(std::move(t)), state_(State::OnlyResult), attached_(1) {}

  template <typename... Args>
  explicit Core(in_place_t, Args&&... args) noexcept(
      std::is_nothrow_constructible<T, Args&&...>::value)
      : result_(in_place, std::forward<Args>(args)...),
        state_(State::OnlyResult),
        attached_(1) {}

  ~Core() {
    DCHECK(attached_ == 0);
    auto state = state_.load(std::memory_order_relaxed);
    switch (state) {
      case State::OnlyResult:
        FOLLY_FALLTHROUGH;

      case State::Done:
        result_.~Result();
        break;

      case State::Proxy:
        proxy_->detachFuture();
        break;

      case State::Empty:
        break;

      default:
        terminate_with<std::logic_error>("~Core unexpected state");
    }
  }

  // Helper class that stores a pointer to the `Core` object and calls
  // `derefCallback` and `detachOne` in the destructor.
  class CoreAndCallbackReference {
   public:
    explicit CoreAndCallbackReference(Core* core) noexcept : core_(core) {}

    ~CoreAndCallbackReference() noexcept {
      detach();
    }

    CoreAndCallbackReference(CoreAndCallbackReference const& o) = delete;
    CoreAndCallbackReference& operator=(CoreAndCallbackReference const& o) =
        delete;
    CoreAndCallbackReference& operator=(CoreAndCallbackReference&&) = delete;

    CoreAndCallbackReference(CoreAndCallbackReference&& o) noexcept
        : core_(exchange(o.core_, nullptr)) {}

    Core* getCore() const noexcept {
      return core_;
    }

   private:
    void detach() noexcept {
      if (core_) {
        core_->derefCallback();
        core_->detachOne();
      }
    }

    Core* core_{nullptr};
  };

  // May be called at most once.
  void doCallback() {
    DCHECK(state_ == State::Done);
    auto x = exchange(executor_, Executor::KeepAlive<>());
    int8_t priority = priority_;

    if (x) {
      exception_wrapper ew;
      // We need to reset `callback_` after it was executed (which can happen
      // through the executor or, if `Executor::add` throws, below). The
      // executor might discard the function without executing it (now or
      // later), in which case `callback_` also needs to be reset.
      // The `Core` has to be kept alive throughout that time, too. Hence we
      // increment `attached_` and `callbackReferences_` by two, and construct
      // exactly two `CoreAndCallbackReference` objects, which call
      // `derefCallback` and `detachOne` in their destructor. One will guard
      // this scope, the other one will guard the lambda passed to the executor.
      attached_.fetch_add(2, std::memory_order_relaxed);
      callbackReferences_.fetch_add(2, std::memory_order_relaxed);
      CoreAndCallbackReference guard_local_scope(this);
      CoreAndCallbackReference guard_lambda(this);
      try {
        auto xPtr = x.get();
        if (LIKELY(x->getNumPriorities() == 1)) {
          xPtr->add([core_ref = std::move(guard_lambda),
                     keepAlive = std::move(x)]() mutable {
            auto cr = std::move(core_ref);
            Core* const core = cr.getCore();
            RequestContextScopeGuard rctx(core->context_);
            core->callback_(std::move(core->result_));
          });
        } else {
          xPtr->addWithPriority(
              [core_ref = std::move(guard_lambda),
               keepAlive = std::move(x)]() mutable {
                auto cr = std::move(core_ref);
                Core* const core = cr.getCore();
                RequestContextScopeGuard rctx(core->context_);
                core->callback_(std::move(core->result_));
              },
              priority);
        }
      } catch (const std::exception& e) {
        ew = exception_wrapper(std::current_exception(), e);
      } catch (...) {
        ew = exception_wrapper(std::current_exception());
      }
      if (ew) {
        RequestContextScopeGuard rctx(context_);
        result_ = Try<T>(std::move(ew));
        callback_(std::move(result_));
      }
    } else {
      attached_.fetch_add(1, std::memory_order_relaxed);
      SCOPE_EXIT {
        context_.~Context();
        callback_.~Callback();
        detachOne();
      };
      RequestContextScopeGuard rctx(context_);
      callback_(std::move(result_));
    }
  }

  void proxyCallback() {
    state_.store(State::Empty, std::memory_order_relaxed);
    proxy_->setExecutor(std::move(executor_), priority_);
    proxy_->setCallback(std::move(callback_), std::move(context_));
    proxy_->detachFuture();
    context_.~Context();
    callback_.~Callback();
  }

  void detachOne() noexcept {
    auto a = attached_.fetch_sub(1, std::memory_order_acq_rel);
    assert(a >= 1);
    if (a == 1) {
      delete this;
    }
  }

  void derefCallback() noexcept {
    auto c = callbackReferences_.fetch_sub(1, std::memory_order_acq_rel);
    assert(c >= 1);
    if (c == 1) {
      context_.~Context();
      callback_.~Callback();
    }
  }

  using Context = std::shared_ptr<RequestContext>;

  union {
    Callback callback_;
  };
  // place result_ next to increase the likelihood that the value will be
  // contained entirely in one cache line
  union {
    Result result_;
    Core* proxy_;
  };
  std::atomic<State> state_;
  std::atomic<unsigned char> attached_;
  std::atomic<unsigned char> callbackReferences_{0};
  std::atomic<bool> interruptHandlerSet_{false};
  SpinLock interruptLock_;
  int8_t priority_{-1};
  Executor::KeepAlive<> executor_;
  union {
    Context context_;
  };
  std::unique_ptr<exception_wrapper> interrupt_{};
  std::function<void(exception_wrapper const&)> interruptHandler_{nullptr};
};

} // namespace detail
} // namespace futures
} // namespace folly