Cleaner disconnect handling. Better fix for issue #71, and also simplifies the…

Cleaner disconnect handling.  Better fix for issue #71, and also simplifies the interface and improves robustness.

c++/src/capnp/capability.c++

@@ -578,9 +578,7 @@ public:
 
   Request<AnyPointer, AnyPointer> newCall(
       uint64_t interfaceId, uint16_t methodId, kj::Maybe<MessageSize> sizeHint) override {
-    auto hook = kj::heap<BrokenRequest>(exception, sizeHint);
-    auto root = hook->message.getRoot<AnyPointer>();
-    return Request<AnyPointer, AnyPointer>(root, kj::mv(hook));
+    return newBrokenRequest(kj::cp(exception), sizeHint);
   }
 
   VoidPromiseAndPipeline call(uint64_t interfaceId, uint16_t methodId,
@@ -626,4 +624,11 @@ kj::Own<PipelineHook> newBrokenPipeline(kj::Exception&& reason) {
   return kj::refcounted<BrokenPipeline>(kj::mv(reason));
 }
 
+Request<AnyPointer, AnyPointer> newBrokenRequest(
+    kj::Exception&& reason, kj::Maybe<MessageSize> sizeHint) {
+  auto hook = kj::heap<BrokenRequest>(kj::mv(reason), sizeHint);
+  auto root = hook->message.getRoot<AnyPointer>();
+  return Request<AnyPointer, AnyPointer>(root, kj::mv(hook));
+}
+
 }  // namespace capnp

c++/src/capnp/capability.h

@@ -419,6 +419,10 @@ kj::Own<ClientHook> newBrokenCap(kj::Exception&& reason);
 kj::Own<PipelineHook> newBrokenPipeline(kj::Exception&& reason);
 // Helper function that creates a pipeline which simply throws exceptions when called.
 
+Request<AnyPointer, AnyPointer> newBrokenRequest(
+    kj::Exception&& reason, kj::Maybe<MessageSize> sizeHint);
+// Helper function that creates a Request object that simply throws exceptions when sent.
+
 // =======================================================================================
 // Extend PointerHelpers for interfaces
 

c++/src/capnp/ez-rpc.c++

@@ -240,7 +240,7 @@ struct EzRpcServer::Impl final: public SturdyRefRestorer<Text>, public kj::TaskS
 
       // Arrange to destroy the server context when all references are gone, or when the
       // EzRpcServer is destroyed (which will destroy the TaskSet).
-      tasks.add(server->network.onDrained().attach(kj::mv(server)));
+      tasks.add(server->network.onDisconnect().attach(kj::mv(server)));
     })));
   }
 

c++/src/capnp/rpc-twoparty-test.c++

@@ -66,7 +66,6 @@ kj::AsyncIoProvider::PipeThread runServer(kj::AsyncIoProvider& ioProvider, int&
     TestRestorer restorer(callCount);
     auto server = makeRpcServer(network, restorer);
     network.onDisconnect().wait(waitScope);
-    network.onDrained().wait(waitScope);
   });
 }
 
@@ -141,9 +140,7 @@ TEST(TwoPartyNetwork, Pipelining) {
   auto rpcClient = makeRpcClient(network);
 
   bool disconnected = false;
-  bool drained = false;
   kj::Promise<void> disconnectPromise = network.onDisconnect().then([&]() { disconnected = true; });
-  kj::Promise<void> drainedPromise = network.onDrained().then([&]() { drained = true; });
 
   {
     // Request the particular capability from the server.
@@ -182,14 +179,12 @@ TEST(TwoPartyNetwork, Pipelining) {
     }
 
     EXPECT_FALSE(disconnected);
-    EXPECT_FALSE(drained);
 
     // What if we disconnect?
     serverThread.pipe->shutdownWrite();
 
     // The other side should also disconnect.
     disconnectPromise.wait(ioContext.waitScope);
-    EXPECT_FALSE(drained);
 
     {
       // Use the now-broken capability.
@@ -213,11 +208,7 @@ TEST(TwoPartyNetwork, Pipelining) {
       EXPECT_EQ(3, callCount);
       EXPECT_EQ(1, reverseCallCount);
     }
-
-    EXPECT_FALSE(drained);
   }
-
-  drainedPromise.wait(ioContext.waitScope);
 }
 
 }  // namespace

c++/src/capnp/rpc-twoparty.c++

@@ -30,37 +30,20 @@ namespace capnp {
 TwoPartyVatNetwork::TwoPartyVatNetwork(kj::AsyncIoStream& stream, rpc::twoparty::Side side,
                                        ReaderOptions receiveOptions)
     : stream(stream), side(side), receiveOptions(receiveOptions), previousWrite(kj::READY_NOW) {
-  {
   auto paf = kj::newPromiseAndFulfiller<void>();
-    drainedPromise = paf.promise.fork();
-    drainedFulfiller.fulfiller = kj::mv(paf.fulfiller);
-  }
-
-  {
-    auto paf = kj::newPromiseAndFulfiller<void>();
-
-    // If the RPC system on this side drops the connection, thus firing onDrained() before
-    // onDisconnected(), we also want to consider ourselves disconnected.  Otherwise, we might
-    // not detect actual disconnect because the RPC system won't attempt to send or receive any
-    // more messages on the connection.  So, we exclusive-join the disconnect promise with the
-    // first branch of drainedPromise.
-    disconnectPromise = paf.promise.exclusiveJoin(drainedPromise.addBranch()).fork();
-
-    disconnectFulfiller = kj::mv(paf.fulfiller);
-  }
+  disconnectPromise = paf.promise.fork();
+  disconnectFulfiller.fulfiller = kj::mv(paf.fulfiller);
 }
 
 void TwoPartyVatNetwork::FulfillerDisposer::disposeImpl(void* pointer) const {
-  KJ_DBG("deref", this, refcount);
   if (--refcount == 0) {
     fulfiller->fulfill();
   }
 }
 
 kj::Own<TwoPartyVatNetworkBase::Connection> TwoPartyVatNetwork::asConnection() {
-  KJ_DBG("ref", &drainedFulfiller, drainedFulfiller.refcount);
-  ++drainedFulfiller.refcount;
-  return kj::Own<TwoPartyVatNetworkBase::Connection>(this, drainedFulfiller);
+  ++disconnectFulfiller.refcount;
+  return kj::Own<TwoPartyVatNetworkBase::Connection>(this, disconnectFulfiller);
 }
 
 kj::Maybe<kj::Own<TwoPartyVatNetworkBase::Connection>> TwoPartyVatNetwork::connectToRefHost(
@@ -102,12 +85,10 @@ public:
 
   void send() override {
     network.previousWrite = network.previousWrite.then([&]() {
-      auto promise = writeMessage(network.stream, message).then([]() {
-        // success; do nothing
-      }, [&](kj::Exception&& exception) {
-        // Exception during write!
-        network.disconnectFulfiller->fulfill();
-      }).eagerlyEvaluate(nullptr);
+      // Note that if the write fails, all further writes will be skipped due to the exception.
+      // We never actually handle this exception because we assume the read end will fail as well
+      // and it's cleaner to handle the failure there.
+      auto promise = writeMessage(network.stream, message).eagerlyEvaluate(nullptr);
       return kj::mv(promise);
     }).attach(kj::addRef(*this));
   }
@@ -145,13 +126,8 @@ kj::Promise<kj::Maybe<kj::Own<IncomingRpcMessage>>> TwoPartyVatNetwork::receiveI
       KJ_IF_MAYBE(m, message) {
         return kj::Own<IncomingRpcMessage>(kj::heap<IncomingMessageImpl>(kj::mv(*m)));
       } else {
-        disconnectFulfiller->fulfill();
         return nullptr;
       }
-    }, [&](kj::Exception&& exception) {
-      disconnectFulfiller->fulfill();
-      kj::throwRecoverableException(kj::mv(exception));
-      return nullptr;
     });
   });
 }

c++/src/capnp/rpc-twoparty.h

@@ -49,19 +49,6 @@ public:
 
   kj::Promise<void> onDisconnect() { return disconnectPromise.addBranch(); }
   // Returns a promise that resolves when the peer disconnects.
-  //
-  // TODO(soon):  Currently this fires when the underlying physical connection breaks.  It should
-  //   fire after the RPC system has detected EOF itself and dropped its connection reference, so
-  //   that it has a chance to reply to connections ended cleanly.
-
-  kj::Promise<void> onDrained() { return drainedPromise.addBranch(); }
-  // Returns a promise that resolves once the peer has disconnected *and* all local objects
-  // referencing this connection have been destroyed.  A caller might use this to decide when it
-  // is safe to destroy the RpcSystem, if it isn't able to reliably destroy all objects using it
-  // directly.
-  //
-  // TODO(soon):  This is not quite designed right.  Those local objects should simply be disabled.
-  //   Their existence should not prevent the RpcSystem from being destroyed.
 
   // implements VatNetwork -----------------------------------------------------
 
@@ -86,14 +73,12 @@ private:
   // second call on the server side.  Never fulfilled, because there is only one connection.
 
   kj::ForkedPromise<void> disconnectPromise = nullptr;
-  kj::Own<kj::PromiseFulfiller<void>> disconnectFulfiller;
-  kj::ForkedPromise<void> drainedPromise = nullptr;
 
   class FulfillerDisposer: public kj::Disposer {
-    // Hack:  TwoPartyVatNetwork is both a VatNetwork and a VatNetwork::Connection.  When all
-    //   references to the Connection have been dropped, then we want onDrained() to fire.  So we
-    //   hand out Own<Connection>s with this disposer attached, so that we can detect when they are
-    //   dropped.
+    // Hack:  TwoPartyVatNetwork is both a VatNetwork and a VatNetwork::Connection.  Whet the RPC
+    //   system detects (or initiates) a disconnection, it drops its reference to the Connection.
+    //   When all references have been dropped, then we want onDrained() to fire.  So we hand out
+    //   Own<Connection>s with this disposer attached, so that we can detect when they are dropped.
 
   public:
     mutable kj::Own<kj::PromiseFulfiller<void>> fulfiller;
@@ -101,7 +86,7 @@ private:
 
     void disposeImpl(void* pointer) const override;
   };
-  FulfillerDisposer drainedFulfiller;
+  FulfillerDisposer disconnectFulfiller;
 
   kj::Own<TwoPartyVatNetworkBase::Connection> asConnection();
   // Returns a pointer to this with the disposer set to drainedFulfiller.