Refactor capability code using fork. Still too much refcounting, though. Maybe…

Refactor capability code using fork.  Still too much refcounting, though.  Maybe this calls for a different design for pipelining...

c++/src/capnp/arena.c++

@@ -113,7 +113,7 @@ public:
   }
 
   VoidPromiseAndPipeline call(uint64_t interfaceId, uint16_t methodId,
-                              CallContextHook& context) const override {
+                              kj::Own<CallContextHook>&& context) const override {
     KJ_FAIL_REQUIRE("Calling capability that was extracted from a message that had no "
                     "capability context.");
   }

c++/src/capnp/capability.h

@@ -342,13 +342,14 @@ public:
   };
 
   virtual VoidPromiseAndPipeline call(uint64_t interfaceId, uint16_t methodId,
-                                      CallContextHook& context) const = 0;
+                                      kj::Own<CallContextHook>&& context) const = 0;
   // Call the object, but the caller controls allocation of the request/response objects.  If the
   // callee insists on allocating this objects itself, it must make a copy.  This version is used
   // when calls come in over the network via an RPC system.  During the call, the context object
-  // may be used from any thread so long as it is only used from one thread at a time.  Once the
-  // returned promise resolves or has been canceled, the context can no longer be used.  The caller
-  // must not allow the ClientHook to be destroyed until the call completes or is canceled.
+  // may be used from any thread so long as it is only used from one thread at a time.  Note that
+  // even if the returned `Promise<void>` is discarded, the call may continue executing if any
+  // pipelined calls are waiting for it; the call is only truly done when the CallContextHook is
+  // destroyed.
   //
   // The call must not begin synchronously, as the caller may hold arbitrary mutexes.
 
@@ -380,10 +381,7 @@ public:
   virtual void allowAsyncCancellation(bool allow) = 0;
   virtual bool isCanceled() = 0;
 
-  virtual Response<ObjectPointer> getResponseForPipeline() = 0;
-  // Get a copy or reference to the response which will be used to execute pipelined calls.  This
-  // will be called no more than once, just after the server implementation successfully returns
-  // from the call.
+  virtual kj::Own<CallContextHook> addRef() = 0;
 };
 
 // =======================================================================================

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

@@ -324,11 +324,11 @@ TEST(Async, Fork) {
 
     auto fork = promise.fork();
 
-    auto branch1 = fork->addBranch().then([](int i) {
+    auto branch1 = fork.addBranch().then([](int i) {
       EXPECT_EQ(123, i);
       return 456;
     });
-    auto branch2 = fork->addBranch().then([](int i) {
+    auto branch2 = fork.addBranch().then([](int i) {
       EXPECT_EQ(123, i);
       return 789;
     });
@@ -360,11 +360,11 @@ TEST(Async, ForkRef) {
 
     auto fork = promise.fork();
 
-    auto branch1 = fork->addBranch().then([](Own<const RefcountedInt>&& i) {
+    auto branch1 = fork.addBranch().then([](Own<const RefcountedInt>&& i) {
       EXPECT_EQ(123, i->i);
       return 456;
     });
-    auto branch2 = fork->addBranch().then([](Own<const RefcountedInt>&& i) {
+    auto branch2 = fork.addBranch().then([](Own<const RefcountedInt>&& i) {
       EXPECT_EQ(123, i->i);
       return 789;
     });

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

@@ -260,7 +260,7 @@ void SimpleEventLoop::wake() const {
 // =======================================================================================
 
 void PromiseBase::absolve() {
-  runCatchingExceptions([this]() { auto deleteMe = kj::mv(node); });
+  runCatchingExceptions([this]() { node = nullptr; });
 }
 
 namespace _ {  // private
@@ -330,9 +330,13 @@ Maybe<const EventLoop&> TransformPromiseNodeBase::getSafeEventLoop() noexcept {
   return loop;
 }
 
+void TransformPromiseNodeBase::dropDependency() {
+  dependency = nullptr;
+}
+
 // -------------------------------------------------------------------
 
-ForkBranchBase::ForkBranchBase(Own<ForkHubBase>&& hubParam): hub(kj::mv(hubParam)) {
+ForkBranchBase::ForkBranchBase(Own<const ForkHubBase>&& hubParam): hub(kj::mv(hubParam)) {
   auto lock = hub->branchList.lockExclusive();
 
   if (lock->lastPtr == nullptr) {
@@ -362,7 +366,7 @@ void ForkBranchBase::hubReady() noexcept {
 
 void ForkBranchBase::releaseHub(ExceptionOrValue& output) {
   KJ_IF_MAYBE(exception, kj::runCatchingExceptions([this]() {
-    auto deleteMe = kj::mv(hub);
+    hub = nullptr;
   })) {
     output.addException(kj::mv(*exception));
   }
@@ -398,7 +402,7 @@ void ForkHubBase::fire() {
     // Dependency is ready.  Fetch its result and then delete the node.
     inner->get(resultRef);
     KJ_IF_MAYBE(exception, kj::runCatchingExceptions([this]() {
-      auto deleteMe = kj::mv(inner);
+      inner = nullptr;
     })) {
       resultRef.addException(kj::mv(*exception));
     }
@@ -525,7 +529,7 @@ void CrossThreadPromiseNodeBase::fire() {
   } else {
     dependency->get(resultRef);
     KJ_IF_MAYBE(exception, kj::runCatchingExceptions([this]() {
-      auto deleteMe = kj::mv(dependency);
+      dependency = nullptr;
     })) {
       resultRef.addException(kj::mv(*exception));
     }

c++/src/kj/async.h

@@ -36,6 +36,8 @@ class SimpleEventLoop;
 template <typename T>
 class Promise;
 template <typename T>
+class ForkedPromise;
+template <typename T>
 class PromiseFulfiller;
 template <typename T>
 struct PromiseFulfillerPair;
@@ -272,6 +274,11 @@ public:
   // Like `Promise::then()`, but schedules the continuation to be executed on *this* EventLoop
   // rather than the thread's current loop.  See Promise::then().
 
+  template <typename T>
+  ForkedPromise<T> fork(Promise<T>&& promise);
+  // Like `Promise::fork()`, but manages the fork on *this* EventLoop rather than the thread's
+  // current loop.  See Promise::fork().
+
   // -----------------------------------------------------------------
   // Low-level interface.
 
@@ -586,14 +593,7 @@ public:
   // After returning, the promise is no longer valid, and cannot be `wait()`ed on or `then()`ed
   // again.
 
-  class Fork {
-  public:
-    virtual Promise<_::Forked<T>> addBranch() = 0;
-    // Add a new branch to the fork.  The branch is equivalent to the original promise, except
-    // that if T is a reference or owned pointer, the target becomes const.
-  };
-
-  Own<Fork> fork();
+  ForkedPromise<T> fork();
   // Forks the promise, so that multiple different clients can independently wait on the result.
   // `T` must be copy-constructable for this to work.  Or, in the special case where `T` is
   // `Own<U>`, `U` must have a method `Own<const U> addRef() const` which returns a new reference
@@ -616,6 +616,27 @@ private:
   friend class _::ForkHub;
 };
 
+template <typename T>
+class ForkedPromise {
+  // The result of `Promise::fork()` and `EventLoop::fork()`.  Allows branches to be created.
+  // Like `Promise<T>`, this is a pass-by-move type.
+
+public:
+  inline ForkedPromise(decltype(nullptr)): hub(nullptr) {}
+
+  Promise<_::Forked<T>> addBranch() const;
+  // Add a new branch to the fork.  The branch is equivalent to the original promise, except
+  // that if T is a reference or owned pointer, the target becomes const.
+
+private:
+  Own<const _::ForkHub<_::FixVoid<T>>> hub;
+
+  inline ForkedPromise(bool, Own<const _::ForkHub<_::FixVoid<T>>>&& hub): hub(kj::mv(hub)) {}
+
+  friend class Promise<T>;
+  friend class EventLoop;
+};
+
 constexpr _::Void READY_NOW = _::Void();
 // Use this when you need a Promise<void> that is already fulfilled -- this value can be implicitly
 // cast to `Promise<void>`.
@@ -881,6 +902,8 @@ private:
   const EventLoop& loop;
   Own<PromiseNode> dependency;
 
+  void dropDependency();
+
   virtual void getImpl(ExceptionOrValue& output) = 0;
 
   template <typename, typename, typename, typename>
@@ -898,6 +921,13 @@ public:
       : TransformPromiseNodeBase(loop, kj::mv(dependency)),
         func(kj::fwd<Func>(func)), errorHandler(kj::fwd<ErrorFunc>(errorHandler)) {}
 
+  ~TransformPromiseNode() noexcept(false) {
+    // We need to make sure the dependency is deleted before we delete the continuations because it
+    // is a common pattern for the continuations to hold ownership of objects that might be in-use
+    // by the dependency.
+    dropDependency();
+  }
+
 private:
   Func func;
   ErrorFunc errorHandler;
@@ -927,7 +957,7 @@ class ForkHubBase;
 
 class ForkBranchBase: public PromiseNode {
 public:
-  ForkBranchBase(Own<ForkHubBase>&& hub);
+  ForkBranchBase(Own<const ForkHubBase>&& hub);
   ~ForkBranchBase();
 
   void hubReady() noexcept;
@@ -946,7 +976,7 @@ protected:
 private:
   EventLoop::Event* onReadyEvent = nullptr;
 
-  Own<ForkHubBase> hub;
+  Own<const ForkHubBase> hub;
   ForkBranchBase* next = nullptr;
   ForkBranchBase** prevPtr = nullptr;
 
@@ -963,7 +993,7 @@ class ForkBranch final: public ForkBranchBase {
   // a const reference.
 
 public:
-  ForkBranch(Own<ForkHubBase>&& hub): ForkBranchBase(kj::mv(hub)) {}
+  ForkBranch(Own<const ForkHubBase>&& hub): ForkBranchBase(kj::mv(hub)) {}
 
   void get(ExceptionOrValue& output) noexcept override {
     const ExceptionOr<T>& hubResult = getHubResultRef().template as<T>();
@@ -1006,7 +1036,7 @@ private:
 };
 
 template <typename T>
-class ForkHub final: public ForkHubBase, public Promise<T>::Fork {
+class ForkHub final: public ForkHubBase {
   // A PromiseNode that implements the hub of a fork.  The first call to Promise::fork() replaces
   // the promise's outer node with a ForkHub, and subsequent calls add branches to that hub (if
   // possible).
@@ -1015,8 +1045,8 @@ public:
   ForkHub(const EventLoop& loop, Own<PromiseNode>&& inner)
       : ForkHubBase(loop, kj::mv(inner), result) {}
 
-  Promise<_::Forked<T>> addBranch() override {
-    return Promise<_::Forked<T>>(false, kj::heap<ForkBranch<T>>(addRef(*this)));
+  Promise<_::Forked<_::UnfixVoid<T>>> addBranch() const {
+    return Promise<_::Forked<_::UnfixVoid<T>>>(false, kj::heap<ForkBranch<T>>(addRef(*this)));
   }
 
 private:
@@ -1261,9 +1291,23 @@ T Promise<T>::wait() {
 }
 
 template <typename T>
-Own<typename Promise<T>::Fork> Promise<T>::fork() {
+ForkedPromise<T> Promise<T>::fork() {
   auto& loop = EventLoop::current();
-  return refcounted<_::ForkHub<T>>(loop, _::makeSafeForLoop<_::FixVoid<T>>(kj::mv(node), loop));
+  return ForkedPromise<T>(false,
+      refcounted<_::ForkHub<_::FixVoid<T>>>(
+          loop, _::makeSafeForLoop<_::FixVoid<T>>(kj::mv(node), loop)));
+}
+
+template <typename T>
+ForkedPromise<T> EventLoop::fork(Promise<T>&& promise) {
+  return ForkedPromise<T>(false,
+      refcounted<_::ForkHub<_::FixVoid<T>>>(*this,
+          _::makeSafeForLoop<_::FixVoid<T>>(kj::mv(promise.node), *this)));
+}
+
+template <typename T>
+Promise<_::Forked<T>> ForkedPromise<T>::addBranch() const {
+  return hub->addBranch();
 }
 
 // =======================================================================================

c++/src/kj/refcount.c++

@@ -22,13 +22,22 @@
 // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 #include "refcount.h"
+#include <memory>
 
 namespace kj {
 
 Refcounted::~Refcounted() noexcept(false) {}
 
 void Refcounted::disposeImpl(void* pointer) const {
-  if (__atomic_sub_fetch(&refcount, 1, __ATOMIC_RELAXED) == 0) {
+  // The load is a fast-path for the common case where this is the last reference.  An acquire-load
+  // is just a regular load on x86.  If there is more than one reference, then we need to do a full
+  // atomic decrement with full memory barrier, because:
+  // - If this is the final decrement then we need to acquire the object state in order to destroy
+  //   it.
+  // - If this is not the final decrement then we need to release the object state so that another
+  //   thread may destroy it.
+  if (__atomic_load_n(&refcount, __ATOMIC_ACQUIRE) == 1 ||
+      __atomic_sub_fetch(&refcount, 1, __ATOMIC_ACQ_REL) == 0) {
     delete this;
   }
 }