Implement fibers on Unix.

Fibers allow code to be written in a synchronous / blocking style while running inside the KJ event loop, by executing the code on an alternate call stack and switching back to the main stack whenever it waits.

We introduce a new function, `kj::startFiber(stackSize, func)`. `func` is executed on the fiber stack. It is passed as its parameter a `WaitScope&`, which can then be passed into the `.wait()` method of any promise in order to wait on the promise in a blocking style. `startFiber()` returns a promise for the eventual value returned by `func()` (much as `evalLater()` and friends do).

This commit implements fibers on Unix via ucontext_t. Windows will come next (and will probably be easier...).

c++/src/kj/async-inl.h

@@ -871,6 +871,77 @@ private:
   }
 };
 
+// -------------------------------------------------------------------
+
+class FiberBase: public PromiseNode, private Event {
+  // Base class for the outer PromiseNode representing a fiber.
+
+public:
+  FiberBase(size_t stackSize, _::ExceptionOrValue& result);
+  ~FiberBase() noexcept(false);
+
+  void start() { armDepthFirst(); }
+  // Call immediately after construction to begin executing the fiber.
+
+  class WaitDoneEvent;
+
+  void onReady(_::Event* event) noexcept override;
+  PromiseNode* getInnerForTrace() override;
+
+protected:
+  bool isFinished() { return state == FINISHED; }
+
+private:
+  enum { WAITING, RUNNING, CANCELED, FINISHED } state;
+
+  size_t stackSize;
+
+  struct Impl;
+  Impl& impl;
+
+  _::PromiseNode* currentInner = nullptr;
+  OnReadyEvent onReadyEvent;
+  _::ExceptionOrValue& result;
+
+  void run();
+  virtual void runImpl(WaitScope& waitScope) = 0;
+
+  struct StartRoutine;
+
+  void switchToFiber();
+  void switchToMain();
+
+  Maybe<Own<Event>> fire() override;
+  // Implements Event. Each time the event is fired, switchToFiber() is called.
+
+  friend class WaitScope;
+  friend void _::waitImpl(Own<_::PromiseNode>&& node, _::ExceptionOrValue& result,
+                          WaitScope& waitScope);
+  friend bool _::pollImpl(_::PromiseNode& node, WaitScope& waitScope);
+};
+
+template <typename Func>
+class Fiber final: public FiberBase {
+public:
+  Fiber(size_t stackSize, Func&& func): FiberBase(stackSize, result), func(kj::fwd<Func>(func)) {}
+
+  typedef FixVoid<decltype(kj::instance<Func&>()(kj::instance<WaitScope&>()))> ResultType;
+
+  void get(ExceptionOrValue& output) noexcept override {
+    KJ_IREQUIRE(isFinished());
+    output.as<ResultType>() = kj::mv(result);
+  }
+
+private:
+  Func func;
+  ExceptionOr<ResultType> result;
+
+  void runImpl(WaitScope& waitScope) override {
+    result.template as<ResultType>() =
+        MaybeVoidCaller<WaitScope&, ResultType>::apply(func, waitScope);
+  }
+};
+
 }  // namespace _ (private)
 
 // =======================================================================================
@@ -1014,6 +1085,17 @@ inline PromiseForResult<Func, void> evalNow(Func&& func) {
   return result;
 }
 
+template <typename Func>
+inline PromiseForResult<Func, WaitScope&> startFiber(size_t stackSize, Func&& func) {
+  typedef _::FixVoid<_::ReturnType<Func, WaitScope&>> ResultT;
+
+  Own<_::FiberBase> intermediate = kj::heap<_::Fiber<Func>>(stackSize, kj::fwd<Func>(func));
+  intermediate->start();
+  auto result = _::PromiseNode::to<_::ChainPromises<_::ReturnType<Func, WaitScope&>>>(
+      _::maybeChain(kj::mv(intermediate), implicitCast<ResultT*>(nullptr)));
+  return _::maybeReduce(kj::mv(result), false);
+}
+
 template <typename T>
 template <typename ErrorFunc>
 void Promise<T>::detach(ErrorFunc&& errorHandler) {

c++/src/kj/async-prelude.h

@@ -188,6 +188,7 @@ class PromiseNode;
 class ChainPromiseNode;
 template <typename T>
 class ForkHub;
+class FiberBase;
 
 class Event;
 class XThreadEvent;

c++/src/kj/async-test.c++

@@ -844,5 +844,93 @@ KJ_TEST("exclusiveJoin both events complete simultaneously") {
   KJ_EXPECT(!joined.poll(waitScope));
 }
 
+KJ_TEST("start a fiber") {
+  EventLoop loop;
+  WaitScope waitScope(loop);
+
+  auto paf = newPromiseAndFulfiller<int>();
+
+  Promise<StringPtr> fiber = startFiber(65536,
+      [promise = kj::mv(paf.promise)](WaitScope& fiberScope) mutable {
+    int i = promise.wait(fiberScope);
+    KJ_EXPECT(i == 123);
+    return "foo"_kj;
+  });
+
+  KJ_EXPECT(!fiber.poll(waitScope));
+
+  paf.fulfiller->fulfill(123);
+
+  KJ_ASSERT(fiber.poll(waitScope));
+  KJ_EXPECT(fiber.wait(waitScope) == "foo");
+}
+
+KJ_TEST("fiber promise chaining") {
+  EventLoop loop;
+  WaitScope waitScope(loop);
+
+  auto paf = newPromiseAndFulfiller<int>();
+  bool ran = false;
+
+  Promise<int> fiber = startFiber(65536,
+      [promise = kj::mv(paf.promise), &ran](WaitScope& fiberScope) mutable {
+    ran = true;
+    return kj::mv(promise);
+  });
+
+  KJ_EXPECT(!ran);
+  KJ_EXPECT(!fiber.poll(waitScope));
+  KJ_EXPECT(ran);
+
+  paf.fulfiller->fulfill(123);
+
+  KJ_ASSERT(fiber.poll(waitScope));
+  KJ_EXPECT(fiber.wait(waitScope) == 123);
+}
+
+KJ_TEST("throw from a fiber") {
+  EventLoop loop;
+  WaitScope waitScope(loop);
+
+  auto paf = newPromiseAndFulfiller<void>();
+
+  Promise<void> fiber = startFiber(65536,
+      [promise = kj::mv(paf.promise)](WaitScope& fiberScope) mutable {
+    promise.wait(fiberScope);
+    KJ_FAIL_EXPECT("wait() should have thrown");
+  });
+
+  KJ_EXPECT(!fiber.poll(waitScope));
+
+  paf.fulfiller->reject(KJ_EXCEPTION(FAILED, "test exception"));
+
+  KJ_ASSERT(fiber.poll(waitScope));
+  KJ_EXPECT_THROW_MESSAGE("test exception", fiber.wait(waitScope));
+}
+
+KJ_TEST("cancel a fiber") {
+  EventLoop loop;
+  WaitScope waitScope(loop);
+
+  auto paf = newPromiseAndFulfiller<int>();
+
+  bool exited = false;
+
+  {
+    Promise<StringPtr> fiber = startFiber(65536,
+        [promise = kj::mv(paf.promise), &exited](WaitScope& fiberScope) mutable {
+      KJ_DEFER(exited = true);
+      int i = promise.wait(fiberScope);
+      KJ_EXPECT(i == 123);
+      return "foo"_kj;
+    });
+
+    KJ_EXPECT(!fiber.poll(waitScope));
+    KJ_EXPECT(!exited);
+  }
+
+  KJ_EXPECT(exited);
+}
+
 }  // namespace
 }  // namespace kj

c++/src/kj/async.c++

@@ -21,6 +21,16 @@
 
 #if _WIN32
 #define WIN32_LEAN_AND_MEAN 1  // lolz
+#elif __APPLE__
+// getcontext() and friends are marked deprecated on MacOS but seemingly no replacement is
+// provided. It appears as if they deprecated it solely because the standards bodies deprecated it,
+// which they seemingly did mainly because the proper sematics are too difficult for them to
+// define. I doubt MacOS would actually remove these functions as they are widely used. But if they
+// do, then I guess we'll need to fall back to using setjmp()/longjmp(), and some sort of hack
+// involving sigaltstack() (and generating a fake signal I guess) in order to initialize the fiber
+// in the first place. Or we could use assembly, I suppose. Either way, ick.
+#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
+#define _XOPEN_SOURCE  // Must be defined to see getcontext() on MacOS.
 #endif
 
 #include "async.h"
@@ -34,6 +44,10 @@
 #include "windows-sanity.h"
 #else
 #include <sched.h>    // just for sched_yield()
+#include <ucontext.h>
+#include <sys/mman.h>
+#include <unistd.h>
+#include <errno.h>
 #endif
 
 #if !KJ_NO_RTTI
@@ -691,6 +705,182 @@ const Executor& getCurrentThreadExecutor() {
   return currentEventLoop().getExecutor();
 }
 
+// =======================================================================================
+// Fiber implementation.
+
+namespace _ {  // private
+
+struct FiberBase::Impl {
+  // This struct serves two purposes:
+  // - It contains OS-specific state that we don't want to declare in the header.
+  // - It is allocated at the top of the fiber's stack area, so the Impl pointer also serves to
+  //   track where the stack was allocated.
+
+  ucontext_t fiberContext;
+  ucontext_t originalContext;
+
+  static Impl& alloc(size_t stackSize) {
+    // TODO(perf): Freelist stacks to avoid TLB flushes.
+
+#ifndef MAP_ANONYMOUS
+#define MAP_ANONYMOUS MAP_ANON
+#endif
+#ifndef MAP_STACK
+#define MAP_STACK 0
+#endif
+
+    size_t pageSize = getPageSize();
+    size_t allocSize = stackSize + pageSize;  // size plus guard page
+
+    // Allocate virtual address space for the stack but make it inaccessible initially.
+    // TODO(someday): Does it make sense to use MAP_GROWSDOWN on Linux? It's a kind of bizarre flag
+    //   that causes the mapping to automatically allocate extra pages (beyond the range specified)
+    //   until it hits something...
+    void* stack = mmap(nullptr, allocSize, PROT_NONE,
+        MAP_ANONYMOUS | MAP_PRIVATE | MAP_STACK, -1, 0);
+    if (stack == MAP_FAILED) {
+      KJ_FAIL_SYSCALL("mmap(new stack)", errno);
+    }
+    KJ_ON_SCOPE_FAILURE({
+      KJ_SYSCALL(munmap(stack, allocSize)) { break; }
+    });
+
+    // Now mark everything except the guard page as read-write. We assume the stack grows down, so
+    // the guard page is at the beginning. No modern architecture uses stacks that grow up.
+    KJ_SYSCALL(mprotect(reinterpret_cast<byte*>(stack) + pageSize,
+                        stackSize, PROT_READ | PROT_WRITE));
+
+    // Stick `Impl` at the top of the stack.
+    Impl& impl = *(reinterpret_cast<Impl*>(reinterpret_cast<byte*>(stack) + allocSize) - 1);
+
+    // Note: mmap() allocates zero'd pages so we don't have to memset() anything here.
+
+    KJ_SYSCALL(getcontext(&impl.fiberContext));
+    impl.fiberContext.uc_stack.ss_size = allocSize - sizeof(Impl);
+    impl.fiberContext.uc_stack.ss_sp = reinterpret_cast<char*>(stack);
+    impl.fiberContext.uc_stack.ss_flags = 0;
+    impl.fiberContext.uc_link = &impl.originalContext;
+
+    return impl;
+  }
+
+  static void free(Impl& impl, size_t stackSize) {
+    size_t allocSize = stackSize + getPageSize();
+    void* stack = reinterpret_cast<byte*>(&impl + 1) - allocSize;
+    KJ_SYSCALL(munmap(stack, allocSize)) { break; }
+  }
+
+  static size_t getPageSize() {
+#ifndef _SC_PAGESIZE
+#define _SC_PAGESIZE _SC_PAGE_SIZE
+#endif
+    static size_t result = sysconf(_SC_PAGE_SIZE);
+    return result;
+  }
+};
+
+struct FiberBase::StartRoutine {
+  static void run(int arg1, int arg2) {
+    // This is the static C-style function we pass to makeContext().
+
+    // POSIX says the arguments are ints, not pointers. So we split our pointer in half in order to
+    // work correctly on 64-bit machines. Gross.
+    uintptr_t ptr = static_cast<uint>(arg1);
+    ptr |= static_cast<uintptr_t>(static_cast<uint>(arg2)) << (sizeof(ptr) * 4);
+    reinterpret_cast<FiberBase*>(ptr)->run();
+  }
+};
+
+FiberBase::FiberBase(size_t stackSizeParam, _::ExceptionOrValue& result)
+    : state(WAITING),
+      // Force stackSize to a reasonable minimum.
+      stackSize(kj::max(stackSizeParam, 65536)),
+      impl(Impl::alloc(stackSize)),
+      result(result) {
+  // Note: Nothing below here can throw. If that changes then we need to call Impl::free(impl)
+  //   on exceptions...
+
+  // POSIX says the arguments are ints, not pointers. So we split our pointer in half in order to
+  // work correctly on 64-bit machines. Gross.
+  uintptr_t ptr = reinterpret_cast<uintptr_t>(this);
+  int arg1 = ptr & ((uintptr_t(1) << (sizeof(ptr) * 4)) - 1);
+  int arg2 = ptr >> (sizeof(ptr) * 4);
+
+  makecontext(&impl.fiberContext, reinterpret_cast<void(*)()>(&StartRoutine::run), 2, arg1, arg2);
+}
+
+FiberBase::~FiberBase() noexcept(false) {
+  KJ_DEFER({
+    Impl::free(impl, stackSize);
+  });
+
+  switch (state) {
+    case WAITING:
+      // We can't just free the stack while the fiber is running. We need to force it to execute
+      // until finished, so we cause it to throw an exception.
+      state = CANCELED;
+      switchToFiber();
+
+      // The fiber should only switch back to the main stack on completion, because any further
+      // calls to wait() would throw before trying to switch.
+      KJ_ASSERT(state == FINISHED);
+      break;
+
+    case RUNNING:
+    case CANCELED:
+      // Bad news.
+      KJ_LOG(FATAL, "fiber tried to destroy itself");
+      ::abort();
+      break;
+
+    case FINISHED:
+      // Normal completion, yay.
+      break;
+  }
+}
+
+Maybe<Own<Event>> FiberBase::fire() {
+  KJ_ASSERT(state == WAITING);
+  state = RUNNING;
+  switchToFiber();
+  return nullptr;
+}
+
+void FiberBase::switchToFiber() {
+  // Switch from the main stack to the fiber. Returns once the fiber either calls switchToMain()
+  // or returns from its main function.
+  KJ_SYSCALL(swapcontext(&impl.originalContext, &impl.fiberContext));
+}
+void FiberBase::switchToMain() {
+  // Switch from the fiber to the main stack. Returns the next time the main stack calls
+  // switchToFiber().
+  KJ_SYSCALL(swapcontext(&impl.fiberContext, &impl.originalContext));
+}
+
+void FiberBase::run() {
+  state = RUNNING;
+  KJ_DEFER(state = FINISHED);
+
+  WaitScope waitScope(currentEventLoop(), *this);
+  KJ_IF_MAYBE(exception, kj::runCatchingExceptions([&]() {
+    runImpl(waitScope);
+  })) {
+    result.addException(kj::mv(*exception));
+  }
+
+  onReadyEvent.arm();
+}
+
+void FiberBase::onReady(_::Event* event) noexcept {
+  onReadyEvent.init(event);
+}
+
+PromiseNode* FiberBase::getInnerForTrace() {
+  return currentInner;
+}
+
+}  // namespace _ (private)
+
 // =======================================================================================
 
 void EventPort::setRunnable(bool runnable) {}
@@ -866,9 +1056,47 @@ void WaitScope::poll() {
 
 namespace _ {  // private
 
+static kj::Exception fiberCanceledException() {
+  // Construct the exception to throw from wait() when the fiber has been canceled (because the
+  // promise returned by startFiber() was dropped before completion).
+  //
+  // TODO(someday): Should we have a dedicated exception type for cancellation? Do we even want
+  //   to build stack traces and such for these?
+  return KJ_EXCEPTION(FAILED, "This fiber is being canceled.");
+};
+
 void waitImpl(Own<_::PromiseNode>&& node, _::ExceptionOrValue& result, WaitScope& waitScope) {
   EventLoop& loop = waitScope.loop;
   KJ_REQUIRE(&loop == threadLocalEventLoop, "WaitScope not valid for this thread.");
+
+  KJ_IF_MAYBE(fiber, waitScope.fiber) {
+    if (fiber->state == FiberBase::CANCELED) {
+      result.addException(fiberCanceledException());
+      return;
+    }
+    KJ_REQUIRE(fiber->state == FiberBase::RUNNING,
+        "This WaitScope can only be used within the fiber that created it.");
+
+    node->setSelfPointer(&node);
+    node->onReady(fiber);
+
+    fiber->currentInner = node;
+    KJ_DEFER(fiber->currentInner = nullptr);
+
+    // Switch to the main stack to run the event loop.
+    fiber->state = FiberBase::WAITING;
+    fiber->switchToMain();
+
+    // The main stack switched back to us, meaning either the event we registered with
+    // node->onReady() fired, or we are being canceled by FiberBase's destructor.
+
+    if (fiber->state == FiberBase::CANCELED) {
+      result.addException(fiberCanceledException());
+      return;
+    }
+
+    KJ_ASSERT(fiber->state == FiberBase::RUNNING);
+  } else {
     KJ_REQUIRE(!loop.running, "wait() is not allowed from within event callbacks.");
 
     BoolEvent doneEvent;
@@ -892,6 +1120,7 @@ void waitImpl(Own<_::PromiseNode>&& node, _::ExceptionOrValue& result, WaitScope
     }
 
     loop.setRunnable(loop.isRunnable());
+  }
 
   node->get(result);
   KJ_IF_MAYBE(exception, kj::runCatchingExceptions([&]() {
@@ -904,6 +1133,7 @@ void waitImpl(Own<_::PromiseNode>&& node, _::ExceptionOrValue& result, WaitScope
 bool pollImpl(_::PromiseNode& node, WaitScope& waitScope) {
   EventLoop& loop = waitScope.loop;
   KJ_REQUIRE(&loop == threadLocalEventLoop, "WaitScope not valid for this thread.");
+  KJ_REQUIRE(waitScope.fiber == nullptr, "poll() is not supported in fibers.");
   KJ_REQUIRE(!loop.running, "poll() is not allowed from within event callbacks.");
 
   BoolEvent doneEvent;

c++/src/kj/async.h

@@ -229,8 +229,16 @@ public:
   //   around them in arbitrary ways.  Therefore, callers really need to know if a function they
   //   are calling might wait(), and the `WaitScope&` parameter makes this clear.
   //
-  // TODO(someday):  Implement fibers, and let them call wait() even when they are handling an
-  //   event.
+  // Usually, there is only one `WaitScope` for each `EventLoop`, and it can only be used at the
+  // top level of the thread owning the loop. Calling `wait()` with this `WaitScope` is what
+  // actually causes the event loop to run at all. This top-level `WaitScope` cannot be used
+  // recursively, so cannot be used within an event callback.
+  //
+  // However, it is possible to obtain a `WaitScope` in lower-level code by using fibers. Use
+  // kj::startFiber() to start some code executing on an alternate call stack. That code will get
+  // its own `WaitScope` allowing it to operate in a synchronous style. In this case, `wait()`
+  // switches back to the main stack in order to run the event loop, returning to the fiber's stack
+  // once the awaited promise resolves.
 
   bool poll(WaitScope& waitScope);
   // Returns true if a call to wait() would complete without blocking, false if it would block.
@@ -244,6 +252,8 @@ public:
   // The first poll() verifies that the promise doesn't resolve early, which would otherwise be
   // hard to do deterministically. The second poll() allows you to check that the promise has
   // resolved and avoid a wait() that might deadlock in the case that it hasn't.
+  //
+  // poll() is not supported in fibers; it will throw an exception.
 
   ForkedPromise<T> fork() KJ_WARN_UNUSED_RESULT;
   // Forks the promise, so that multiple different clients can independently wait on the result.
@@ -380,6 +390,29 @@ PromiseForResult<Func, void> evalLast(Func&& func) KJ_WARN_UNUSED_RESULT;
 // callback enqueues new events, then latter callbacks will not execute until those events are
 // drained.
 
+template <typename Func>
+PromiseForResult<Func, WaitScope&> startFiber(size_t stackSize, Func&& func) KJ_WARN_UNUSED_RESULT;
+// Executes `func()` in a fiber, returning a promise for the eventual reseult. `func()` will be
+// passed a `WaitScope&` as its parameter, allowing it to call `.wait()` on promises. Thus, `func()`
+// can be written in a synchronous, blocking style, instead of using `.then()`. This is often much
+// easier to write and read, and may even be significantly faster if it allows the use of stack
+// allocation rather than heap allocation.
+//
+// However, fibers have a major disadvantage: memory must be allocated for the fiber's call stack.
+// The entire stack must be allocated at once, making it necessary to choose a stack size upfront
+// that is big enough for whatever the fiber needs to do. Estimating this is often difficult. That
+// said, over-estimating is not too terrible since pages of the stack will actually be allocated
+// lazily when first accessed; actual memory usage will correspond to the "high watermark" of the
+// actual stack usage. That said, this lazy allocation forces page faults, which can be quite slow.
+// Worse, freeing a stack forces a TLB flush and shootdown -- all currently-executing threads will
+// have to be interrupted to flush their CPU cores' TLB caches.
+//
+// In short, when performance matters, you should try to avoid creating fibers very frequently.
+//
+// TODO(perf): We should add a mechanism for freelisting stacks. However, this improves CPU usage
+//   at the expense of memory usage: stacks on the freelist will consume however many pages they
+//   used at their high watermark, forever.
+
 template <typename T>
 Promise<Array<T>> joinPromises(Array<Promise<T>>&& promises);
 // Join an array of promises into a promise for an array.
@@ -903,20 +936,31 @@ class WaitScope {
 
 public:
   inline explicit WaitScope(EventLoop& loop): loop(loop) { loop.enterScope(); }
-  inline ~WaitScope() { loop.leaveScope(); }
+  inline ~WaitScope() { if (fiber == nullptr) loop.leaveScope(); }
   KJ_DISALLOW_COPY(WaitScope);
 
   void poll();
   // Pumps the event queue and polls for I/O until there's nothing left to do (without blocking).
+  //
+  // Not supported in fibers.
 
   void setBusyPollInterval(uint count) { busyPollInterval = count; }
   // Set the maximum number of events to run in a row before calling poll() on the EventPort to
   // check for new I/O.
+  //
+  // This has no effect when used in a fiber.
 
 private:
   EventLoop& loop;
   uint busyPollInterval = kj::maxValue;
+
+  kj::Maybe<_::FiberBase&> fiber;
+
+  explicit WaitScope(EventLoop& loop, _::FiberBase& fiber)
+      : loop(loop), fiber(fiber) {}
+
   friend class EventLoop;
+  friend class _::FiberBase;
   friend void _::waitImpl(Own<_::PromiseNode>&& node, _::ExceptionOrValue& result,
                           WaitScope& waitScope);
   friend bool _::pollImpl(_::PromiseNode& node, WaitScope& waitScope);