yield.hpp

//          Copyright Oliver Kowalke, Nat Goodspeed 2015.
// Distributed under the Boost Software License, Version 1.0.
//    (See accompanying file LICENSE_1_0.txt or copy at
//          http://www.boost.org/LICENSE_1_0.txt)

#pragma once

#include <boost/asio/async_result.hpp>
#include <boost/asio/detail/config.hpp>
#include <boost/assert.hpp>
#include <boost/atomic.hpp>
#include <boost/intrusive_ptr.hpp>
#include <boost/system/error_code.hpp>
#include <boost/system/system_error.hpp>
#include <boost/throw_exception.hpp>

#include <boost/fiber/mutex.hpp>

#include <mutex>  // std::unique_lock

namespace util {
namespace fibers_ext {
namespace detail {

namespace fbs = ::boost::fibers;

//[fibers_asio_yield_completion
// Bundle a completion bool flag with a spinlock to protect it.
struct yield_completion {
  enum state_t { init, waiting, complete };

  typedef fbs::detail::spinlock mutex_t;
  typedef std::unique_lock<mutex_t> lock_t;
  typedef boost::intrusive_ptr<yield_completion> ptr_t;

  std::atomic<std::size_t> use_count_{0};
  mutex_t mtx_{};
  state_t state_{init};

  void wait() {
    // yield_handler_base::operator()() will set state_ `complete` and
    // attempt to wake a suspended fiber. It would be Bad if that call
    // happened between our detecting (complete != state_) and suspending.
    lock_t lk{mtx_};
    // If state_ is already set, we're done here: don't suspend.
    if (complete != state_) {
      state_ = waiting;
      // suspend(unique_lock<spinlock>) unlocks the lock in the act of
      // resuming another fiber
      fbs::context::active()->suspend(lk);
    }
  }

  friend void intrusive_ptr_add_ref(yield_completion* yc) noexcept {
    BOOST_ASSERT(nullptr != yc);
    yc->use_count_.fetch_add(1, std::memory_order_relaxed);
  }

  friend void intrusive_ptr_release(yield_completion* yc) noexcept {
    BOOST_ASSERT(nullptr != yc);
    if (1 == yc->use_count_.fetch_sub(1, std::memory_order_release)) {
      std::atomic_thread_fence(std::memory_order_acquire);
      delete yc;
    }
  }
};
//]

//[fibers_asio_yield_handler_base
// This class encapsulates common elements between yield_handler<T> (capturing
// a value to return from asio async function) and yield_handler<void> (no
// such value). See yield_handler<T> and its <void> specialization below. Both
// yield_handler<T> and yield_handler<void> are passed by value through
// various layers of asio functions. In other words, they're potentially
// copied multiple times. So key data such as the yield_completion instance
// must be stored in our async_result<yield_handler<>> specialization, which
// should be instantiated only once.
class yield_handler_base {
 public:
  yield_handler_base(yield_t const& y)
      :  // capture the context* associated with the running fiber
        ctx_{fbs::context::active()},
        // capture the passed yield_t
        yt_(y) {}

  // completion callback passing only (error_code)
  void operator()(boost::system::error_code const& ec) {
    BOOST_ASSERT_MSG(ycomp_,
                     "Must inject yield_completion* "
                     "before calling yield_handler_base::operator()()");
    BOOST_ASSERT_MSG(yt_.ec_,
                     "Must inject boost::system::error_code* "
                     "before calling yield_handler_base::operator()()");
    // If originating fiber is busy testing state_ flag, wait until it
    // has observed (completed != state_).
    yield_completion::lock_t lk{ycomp_->mtx_};
    yield_completion::state_t state = ycomp_->state_;
    // Notify a subsequent yield_completion::wait() call that it need not
    // suspend.
    ycomp_->state_ = yield_completion::complete;
    // set the error_code bound by yield_t
    *yt_.ec_ = ec;
    // unlock the lock that protects state_
    lk.unlock();
    // If ctx_ is still active, e.g. because the async operation
    // immediately called its callback (this method!) before the asio
    // async function called async_result_base::get(), we must not set it
    // ready.
    if (yield_completion::waiting == state) {
      // wake the fiber
      fbs::context::active()->schedule(ctx_);
    }
  }

  void bind(yield_completion::ptr_t ptr, boost::system::error_code* ec) {
    // if yield_t didn't bind an error_code, make yield_handler_base's
    // error_code* point to an error_code local to this object so
    // yield_handler_base::operator() can unconditionally store through
    // its error_code*
    ycomp_ = std::move(ptr);
    if (!yt_.ec_) {
      yt_.ec_ = ec;
    }
  }

 private:
  boost::fibers::context* ctx_;
  yield_t yt_;
  // We depend on this pointer to yield_completion, which will be injected
  // by async_result.
  yield_completion::ptr_t ycomp_{};
};
//]

//[fibers_asio_yield_handler_T
// asio uses handler_type<completion token type, signature>::type to decide
// what to instantiate as the actual handler. Below, we specialize
// handler_type< yield_t, ... > to indicate yield_handler<>. So when you pass
// an instance of yield_t as an asio completion token, asio selects
// yield_handler<> as the actual handler class.
template <typename T> class yield_handler : public yield_handler_base {
  friend class ::boost::asio::async_result<util::fibers_ext::yield_t,
                                           void(boost::system::error_code, T)>;

 public:
  // asio passes the completion token to the handler constructor
  explicit yield_handler(yield_t const& y) : yield_handler_base{y} {}

  // completion callback passing only value (T)
  void operator()(T t) {
    // just like callback passing success error_code
    (*this)(boost::system::error_code(), std::move(t));
  }

  // completion callback passing (error_code, T)
  void operator()(boost::system::error_code const& ec, T t) {
    BOOST_ASSERT_MSG(value_,
                     "Must inject value ptr "
                     "before caling yield_handler<T>::operator()()");
    // move the value to async_result<> instance BEFORE waking up a
    // suspended fiber
    *value_ = std::move(t);
    // forward the call to base-class completion handler
    yield_handler_base::operator()(ec);
  }

 private:
  // pointer to destination for eventual value
  // this must be injected by async_result before operator()() is called
  T* value_{nullptr};
};
//]

//[fibers_asio_yield_handler_void
// yield_handler<void> is like yield_handler<T> without value_. In fact it's
// just like yield_handler_base.
template <> class yield_handler<void> : public yield_handler_base {
 public:
  explicit yield_handler(yield_t const& y) : yield_handler_base{y} {}

  // nullary completion callback
  void operator()() { (*this)(boost::system::error_code()); }

  // inherit operator()(error_code) overload from base class
  using yield_handler_base::operator();
};
//]

// Specialize asio_handler_invoke hook to ensure that any exceptions thrown
// from the handler are propagated back to the caller
// template <typename Fn, typename T> void asio_handler_invoke(Fn&& fn, yield_handler<T>*) { fn(); }

//[fibers_asio_async_result_base
// Factor out commonality between async_result<yield_handler<T>> and
// async_result<yield_handler<void>>
class async_result_base {
 public:
  explicit async_result_base(yield_handler_base& h) : ycomp_{new yield_completion{}} {
    // Inject ptr to our yield_completion instance into this
    // yield_handler<>.
    h.bind(ycomp_, &ec_);
  }

  void get() {
    // Unless yield_handler_base::operator() has already been called,
    // suspend the calling fiber until that call.
    ycomp_->wait();
    // The only way our own ec_ member could have a non-default value is
    // if our yield_handler did not have a bound error_code AND the
    // completion callback passed a non-default error_code.
    if (ec_) {
      boost::throw_exception(boost::system::system_error{ec_});
    }
  }

 private:
  // If yield_t does not bind an error_code instance, store into here.
  boost::system::error_code ec_{};
  yield_completion::ptr_t ycomp_;
};
//]

}  // namespace detail
}  // namespace fibers_ext
}  // namespace util

namespace boost {
namespace asio {

//[fibers_asio_async_result_T
// asio constructs an async_result<> instance from the yield_handler specified
// by handler_type<>::type. A particular asio async method constructs the
// yield_handler, constructs this async_result specialization from it, then
// returns the result of calling its get() method.
template <typename T>
class async_result<util::fibers_ext::yield_t, void(boost::system::error_code, T)>
    : public util::fibers_ext::detail::async_result_base {
 public:
  // type returned by get()
  typedef T return_type;
  typedef util::fibers_ext::detail::yield_handler<T> completion_handler_type;

  explicit async_result(completion_handler_type& h)
      : util::fibers_ext::detail::async_result_base{h} {
    // Inject ptr to our value_ member into yield_handler<>: result will
    // be stored here.
    h.value_ = &value_;
  }

  // asio async method returns result of calling get()
  return_type get() {
    util::fibers_ext::detail::async_result_base::get();
    return std::move(value_);
  }

 private:
  return_type value_{};
};
//]

//[fibers_asio_async_result_void
// Without the need to handle a passed value, our yield_handler<void>
// specialization is just like async_result_base.
template <>
class async_result<util::fibers_ext::yield_t, void(boost::system::error_code)>
    : public util::fibers_ext::detail::async_result_base {
 public:
  typedef void return_type;
  typedef util::fibers_ext::detail::yield_handler<void> completion_handler_type;

  explicit async_result(completion_handler_type& h)
      : util::fibers_ext::detail::async_result_base{h} {}
};
//]

}  // namespace asio
}  // namespace boost