From 42da9f41019a2625c61c1dec58c5567a24c71a79 Mon Sep 17 00:00:00 2001
From: Jason Paryani <github@jparyani.com>
Date: Mon, 18 Nov 2013 11:58:03 -0800
Subject: [PATCH] Change `run_forever` to use onDisconnect

---
 capnp/capnp.pyx     | 3 +--
 capnp/capnp_cpp.pxd | 2 ++
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/capnp/capnp.pyx b/capnp/capnp.pyx
index 8c38762..0d1b61e 100644
--- a/capnp/capnp.pyx
+++ b/capnp/capnp.pyx
@@ -1582,8 +1582,7 @@ cdef class RpcServer:
         del self.thisptr
 
     def run_forever(self):
-        p = PromiseFulfillerPair()
-        self.loop.wait(p)
+        self.loop.wait(_VoidPromise()._init(deref(self.network.thisptr).onDisconnect()))
 
     # TODO: add restore functionality here?
 
diff --git a/capnp/capnp_cpp.pxd b/capnp/capnp_cpp.pxd
index 3b521d3..f780e95 100644
--- a/capnp/capnp_cpp.pxd
+++ b/capnp/capnp_cpp.pxd
@@ -204,6 +204,8 @@ cdef extern from "capnp/rpc-twoparty.h" namespace " ::capnp":
         SERVER" ::capnp::rpc::twoparty::Side::SERVER"
     cdef cppclass TwoPartyVatNetwork:
         TwoPartyVatNetwork(EventLoop &, AsyncIoStream& stream, Side)
+        VoidPromise onDisconnect()
+        VoidPromise onDrain()
     RpcSystem makeRpcServer(TwoPartyVatNetwork&, PyRestorer&, EventLoop&)
     RpcSystem makeRpcClient(TwoPartyVatNetwork&, EventLoop&)