proxy-io.h: Add Connection disconnect and waitDrained methods

example/calculator.cpp

@@ -10,7 +10,6 @@
 #include <fstream>
 #include <functional>
 #include <iostream>
-#include <kj/async.h>
 #include <kj/common.h>
 #include <kj/memory.h>
 #include <memory>

example/printer.cpp

@@ -10,7 +10,6 @@
 #include <cstring> // IWYU pragma: keep
 #include <fstream>
 #include <iostream>
-#include <kj/async.h>
 #include <kj/common.h>
 #include <kj/memory.h>
 #include <memory>

include/mp/proxy.h

@@ -20,6 +20,7 @@
 namespace mp {
 class Connection;
 class EventLoop;
+struct ServerObjectTracker;
 //! Mapping from capnp interface type to proxy client implementation (specializations are generated by
 //! proxy-codegen.cpp).
 template <typename Interface> struct ProxyClient; // IWYU pragma: export
@@ -172,6 +173,12 @@ struct ProxyServerBase : public virtual Interface_::Server
      * wrapped. */
     std::shared_ptr<Impl> m_impl;
     ProxyContext m_context;
+    //! Live-object tracker shared with this object's Connection, incremented
+    //! in the constructor and decremented in the destructor so shutdown code
+    //! can wait for a disconnected connection's server objects to drain. Held
+    //! via shared_ptr so it remains valid if this object (kept alive by an
+    //! in-flight call) outlives the Connection. See ServerObjectTracker.
+    std::shared_ptr<ServerObjectTracker> m_server_objects;
 };
 
 //! Customizable (through template specialization) base class which ProxyServer

include/mp/type-context.h

@@ -179,10 +179,10 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                 if (erase_thread) {
                     // Look up the thread again without using existing
                     // iterator since entry may no longer be there after
-                    // a disconnect. Destroy node after releasing
-                    // Waiter::m_mutex, so the ProxyClient<Thread>
-                    // destructor is able to use EventLoop::mutex
-                    // without violating lock order.
+                    // a disconnect. Destroy the node after releasing
+                    // Waiter::m_mutex, so ~ProxyClient<Thread> does not
+                    // run with the mutex held (see the SetThread
+                    // disconnect callback).
                     ConnThreads::node_type removed;
                     {
                         Lock lock(thread_context.waiter->m_mutex);

src/mp/proxy.cpp

@@ -13,6 +13,7 @@
 #include <atomic>
 #include <capnp/capability.h>
 #include <capnp/common.h> // IWYU pragma: keep
+#include <capnp/rpc-twoparty.h>
 #include <capnp/rpc.h>
 #include <condition_variable>
 #include <functional>
@@ -111,9 +112,32 @@ Connection::~Connection() noexcept(false)
     // Connection destructor is always called on the event loop thread. If this
     // is a local disconnect, it will trigger I/O, so this needs to run on the
     // event loop thread, and if there was a remote disconnect, this is called
-    // by an onDisconnect callback directly from the event loop thread.
+    // by a TwoPartyVatNetwork::onDisconnect callback directly from the event
+    // loop thread.
+    assert(std::this_thread::get_id() == m_loop->m_thread_id);
+    disconnect();
+}
+
+void Connection::disconnect()
+{
+    // Disconnecting triggers I/O and tears down capnp state, so it must run on
+    // the event loop thread, like the destructor.
     assert(std::this_thread::get_id() == m_loop->m_thread_id);
 
+    // m_network is reset at the end of teardown below, so treat it being null
+    // as the "already disconnected" state: a second call (including the one
+    // from the destructor) is a no-op.
+    if (!m_network) return;
+
+    // Expire m_alive before severing the connection below so onDisconnect
+    // handlers will not trigger and delete this Connection object. The
+    // onDisconnect handlers trigger on remote disconnects and automatically
+    // delete Connection objects. But on local disconnects, they should not
+    // trigger, because local code that disconnects is responsible for freeing
+    // Connection objects, and it may want to wait for in-flight calls to finish
+    // before destroying them.
+    m_alive.reset();
+
     // Try to cancel any calls that may be executing.
     m_canceler.cancel("Interrupted by disconnect");
 
@@ -132,13 +156,13 @@ Connection::~Connection() noexcept(false)
     //
     // Sending pending data is important if the connection is a socketpair
     // because when one side of the socketpair is closed, the other side doesn't
-    // seem to receive any onDisconnect event. So it is important for the other
-    // side to instead receive Cap'n Proto "release" messages (see `struct
-    // Release` in capnp/rpc.capnp) from local Client objects being destroyed so
-    // the remote side can free resources and shut down cleanly. Without this,
-    // when one side of a socket pair is closed the other side may not receive
-    // these messages, preventing the remote side from freeing ProxyServer
-    // resources and shutting down cleanly.
+    // seem to receive any TwoPartyVatNetwork::onDisconnect event. So it is
+    // important for the other side to instead receive Cap'n Proto "release"
+    // messages (see `struct Release` in capnp/rpc.capnp) from local Client
+    // objects being destroyed so the remote side can free resources and shut
+    // down cleanly. Without this, when one side of a socket pair is closed the
+    // other side may not receive these messages, preventing the remote side
+    // from freeing ProxyServer resources and shutting down cleanly.
     // Use kj::runCatchingExceptions instead of try/catch because on macOS with
     // dynamic libraries, kj::Exception typeinfo differs between libcapnp and
     // the calling binary, so catch (const kj::Exception&) silently fails to
@@ -184,28 +208,49 @@ Connection::~Connection() noexcept(false)
     // connection implementing the Init interface and handling the Init.makeX() calls.
     //
     // Either way when a connection is closed, capnp behavior is to call all
-    // ProxyServer object destructors first, and then trigger an onDisconnect
-    // callback.
+    // ProxyServer object destructors first, and then trigger a
+    // TwoPartyVatNetwork::onDisconnect callback.
     //
-    // On incoming side of the connection, the onDisconnect callback is written
-    // to delete the Connection object from the m_incoming_connections and call
-    // this destructor which calls Connection::disconnect.
+    // On incoming side of the connection, the TwoPartyVatNetwork::onDisconnect
+    // callback is written to delete the Connection object from the
+    // m_incoming_connections list and call this destructor.
     //
-    // On the outgoing side, the Connection object is owned by top level client
-    // object client, which onDisconnect handler doesn't have ready access to,
-    // so onDisconnect handler just calls Connection::disconnect directly
-    // instead.
+    // On the outgoing side, the Connection object is heap-allocated and owned
+    // by a top-level ProxyClient object. In this case, the
+    // TwoPartyVatNetwork::onDisconnect handler deletes the Connection object
+    // directly, calling this destructor, and loop below sets Connection
+    // pointers in all associated ProxyClient objects to null.
     //
     // Either way disconnect code runs in the event loop thread and called both
     // on clean and unclean shutdowns. In unclean shutdown case when the
     // connection is broken, sync and async cleanup lists will be filled with
     // callbacks. In the clean shutdown case both lists will be empty.
-    Lock lock{m_loop->m_mutex};
-    while (!m_sync_cleanup_fns.empty()) {
-        CleanupList fn;
-        fn.splice(fn.begin(), m_sync_cleanup_fns, m_sync_cleanup_fns.begin());
-        Unlock(lock, fn.front());
+    {
+        Lock lock{m_loop->m_mutex};
+        while (!m_sync_cleanup_fns.empty()) {
+            CleanupList fn;
+            fn.splice(fn.begin(), m_sync_cleanup_fns, m_sync_cleanup_fns.begin());
+            Unlock(lock, fn.front());
+        }
     }
+
+    // Release Thread capabilities so idle worker threads are stopped and
+    // joined at disconnect time, whether or not this object is destroyed right
+    // away. (A worker thread currently executing a call body is unaffected:
+    // its ProxyServer<Thread> object is pinned by the post() call and released
+    // when the body finishes.)
+    m_thread_pool.clear();
+    m_thread_map = nullptr;
+
+    // Destroy the network and close the stream so the peer observes the
+    // disconnect, reading EOF and failing its outstanding calls with
+    // DISCONNECTED errors. This has to be explicit because when disconnect()
+    // is called without destroying this object, nothing else severs the
+    // transport: m_rpc_system.reset() above stops reading from the stream but
+    // does not reliably close it. The network is destroyed first since it
+    // references the stream.
+    m_network.reset();
+    m_stream = nullptr;
 }
 
 CleanupIt Connection::addSyncCleanup(std::function<void()> fn)
@@ -218,7 +263,8 @@ CleanupIt Connection::addSyncCleanup(std::function<void()> fn)
     // order should not be significant because the cleanup callbacks run
     // synchronously in a single batch when the connection is broken, and they
     // only reset the connection pointers in the client objects without actually
-    // deleting the client objects.
+    // deleting the client objects, or update ListenConnections max_connections
+    // bookkeeping.
     return m_sync_cleanup_fns.emplace(m_sync_cleanup_fns.begin(), std::move(fn));
 }
 
@@ -394,27 +440,36 @@ std::tuple<ConnThread, bool> SetThread(GuardedRef<ConnThreads> threads, Connecti
     }
     if (inserted) {
         thread->second.emplace(make_thread(), connection, /* destroy_connection= */ false);
-        thread->second->m_disconnect_cb = connection->addSyncCleanup([threads, thread] {
-            // Note: it is safe to use the `thread` iterator in this cleanup
-            // function, because the iterator would only be invalid if the map entry
-            // was removed, and if the map entry is removed the ProxyClient<Thread>
-            // destructor unregisters the cleanup.
-
-            // Connection is being destroyed before thread client is, so reset
-            // thread client m_disconnect_cb member so thread client destructor does not
-            // try to unregister this callback after connection is destroyed.
-            thread->second->m_disconnect_cb.reset();
-
-            // Remove connection pointer about to be destroyed from the map
-            const Lock lock(threads.mutex);
-            threads.ref.erase(thread);
+        thread->second->m_disconnect_cb = connection->addSyncCleanup([threads, connection] {
+            // Remove and destroy this connection's map entry, unless the
+            // thread owning the map is exiting and ~ThreadContext already took
+            // the entry, in which case that thread destroys it. Look the entry
+            // up by key rather than capturing the iterator, which would be
+            // invalid in that case.
+            ConnThreads::node_type removed;
+            {
+                const Lock lock(threads.mutex);
+                auto it = threads.ref.find(connection);
+                if (it == threads.ref.end()) return;
+                // Reset so ~ProxyClient<Thread> does not try to unregister
+                // this callback, which is already running.
+                it->second->m_disconnect_cb.reset();
+                removed = threads.ref.extract(it);
+            }
+            // The entry is destroyed here with Waiter::m_mutex released, as in
+            // ~ThreadContext. (Holding it would be safe on the event loop
+            // thread, where EventLoop::sync() runs directly, but releasing it
+            // keeps the two consistent.)
         });
     }
     return {thread, inserted};
 }
 
 ProxyClient<Thread>::~ProxyClient()
 {
+    EventLoop& loop{*m_context.loop};
+    if (loop.testing_hook_thread_client_destroy) loop.testing_hook_thread_client_destroy(this);
+
     // If thread is being destroyed before connection is destroyed, remove the
     // cleanup callback that was registered to handle the connection being
     // destroyed before the thread being destroyed.
@@ -424,13 +479,43 @@ ProxyClient<Thread>::~ProxyClient()
         // between this thread trying to remove the callback and the disconnect
         // handler attempting to call it.
         m_context.loop->sync([&]() {
-            if (m_disconnect_cb) {
+            // Skip if the connection was disconnected while this thread waited
+            // for the event loop: Connection::disconnect() has already run and
+            // freed the cleanup list m_disconnect_cb points into, and the
+            // ProxyClientBase disconnect callback has nulled
+            // m_context.connection. (The SetThread callback resets
+            // m_disconnect_cb only when it finds this object in the thread
+            // map, so if ~ThreadContext took the entry first, m_disconnect_cb
+            // is still set here.)
+            if (m_disconnect_cb && m_context.connection) {
                 m_context.connection->removeSyncCleanup(*m_disconnect_cb);
             }
         });
     }
 }
 
+ThreadContext::~ThreadContext()
+{
+    // Server threads created by ProxyServer<ThreadMap>::makeThread have no
+    // waiter here: ~ProxyServer<Thread> moves it away and clears the maps
+    // before the thread exits.
+    if (!waiter) return;
+
+    // Take the thread client maps under Waiter::m_mutex, because the SetThread
+    // disconnect callback removes entries from them concurrently on the event
+    // loop thread when connections are broken, and whichever side removes an
+    // entry destroys it. Destroy the maps with the mutex released, because
+    // ~ProxyClient<Thread> blocks in EventLoop::sync() and the event loop
+    // thread locks Waiter::m_mutex itself (see the blocking rule in the
+    // Waiter::m_mutex documentation).
+    ConnThreads request, callback;
+    {
+        const Lock lock(waiter->m_mutex);
+        request.swap(request_threads);
+        callback.swap(callback_threads);
+    }
+}
+
 ProxyServer<Thread>::ProxyServer(Connection& connection, ThreadContext& thread_context, std::thread&& thread)
     : m_loop{*connection.m_loop}, m_thread_context(thread_context), m_thread(std::move(thread))
 {

test/mp/test/connect_tests.cpp

@@ -118,11 +118,11 @@ KJ_TEST("ConnectStream defers disconnect failure to the first IPC request for in
     } catch (const std::runtime_error& e) {
         std::string_view reason = e.what();
 
-        // There is a race between the event loop detecting the disconnect
-        // and foo->add() being called. If the onDisconnect callback fires
-        // first and nulls m_context.connection, the error is "called after
-        // disconnect"; if foo->add() submits before the callback fires,
-        // the error is "interrupted by disconnect".
+        // There is a race between the event loop detecting the disconnect and
+        // foo->add() being called. If the TwoPartyVatNetwork::onDisconnect
+        // callback fires first and nulls m_context.connection, the error is
+        // "called after disconnect"; if foo->add() submits before the callback
+        // fires, the error is "interrupted by disconnect".
         KJ_EXPECT(reason == "IPC client method called after disconnect." ||
                   reason == "IPC client method call interrupted by disconnect.");
     }

test/mp/test/listen_tests.cpp

@@ -18,6 +18,7 @@
 #include <kj/debug.h>
 #include <kj/memory.h>
 #include <kj/test.h>
+#include <list>
 #include <memory>
 #include <mp/proxy.h>
 #include <mp/proxy-io.h>
@@ -192,6 +193,34 @@ KJ_TEST("ListenConnections enforces a local connection limit")
     KJ_EXPECT(client3->client->add(3, 4) == 7);
 }
 
+KJ_TEST("ListenConnections resumes after a local disconnect")
+{
+    // A connection closed locally (here by erasing it from the incoming-
+    // connection list) must still free the listener's slot so the listener
+    // resumes accepting. The counter decrement runs as an addSyncCleanup cleanup;
+    // if it were in the onDisconnect handler it would be canceled by the
+    // local close and the slot would stay stuck.
+    ListenSetup server(/*max_connections=*/1);
+
+    auto client1 = std::make_unique<ClientSetup>(server.listener.MakeConnectedSocket());
+    server.WaitForConnectedCount(1);
+    KJ_EXPECT(client1->client->add(1, 2) == 3);
+
+    auto client2 = std::make_unique<ClientSetup>(server.listener.MakeConnectedSocket());
+    (**server.m_loop_ref).sync([] {});
+    KJ_EXPECT(server.ConnectedCount() == 1);
+
+    // Close the first connection locally on the event loop thread.
+    EventLoop& loop{**server.m_loop_ref};
+    loop.sync([&] {
+        KJ_REQUIRE(loop.m_incoming_connections.size() == 1);
+        loop.m_incoming_connections.pop_front();
+    });
+
+    server.WaitForConnectedCount(2);
+    KJ_EXPECT(client2->client->add(2, 3) == 5);
+}
+
 KJ_TEST("ListenConnections accepts multiple connections")
 {
     // With max-connections=2, two clients should be accepted and usable at the