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

include/mp/proxy-io.h

@@ -22,6 +22,7 @@
 #include <map>
 #include <memory>
 #include <optional>
+#include <ranges>
 #include <sstream>
 #include <string>
 #include <thread>
@@ -298,6 +299,9 @@ class EventLoop
     //! Check if loop should exit.
     bool done() const MP_REQUIRES(m_mutex);
 
+    //! View of incoming connections yielding Connection& for each entry.
+    auto incoming_connections() { return std::views::all(m_incoming_connections); }
+
     //! Process name included in thread names so combined debug output from
     //! multiple processes is easier to understand.
     const char* m_exe_name;
@@ -432,6 +436,72 @@ struct Waiter
     std::optional<kj::Function<void()>> m_fn MP_GUARDED_BY(m_mutex);
 };
 
+//! Counter tracking the number of live ProxyServer objects associated with a
+//! Connection, used to wait for a disconnected connection's server side to
+//! become quiescent (see Connection::waitDrained).
+//!
+//! Why counting live server objects is a valid "no server call body running"
+//! signal: a ProxyServer object is reference counted and is not destroyed
+//! until its outstanding calls finish. Cap'n Proto keeps the target capability
+//! alive for the duration of a call, and the mp.Context PassField overload and
+//! ProxyServer<Thread>::post() additionally pin it (self = thisCap()) until
+//! the call body running on a worker thread completes and its result is
+//! delivered. So "object destroyed" implies "its call bodies finished", and a
+//! connection whose live-object count reached zero after a disconnect has no
+//! server code running. This matters because disconnecting only cancels the
+//! KJ promise of an in-flight call; it does not interrupt a call body that
+//! was already dispatched to a worker thread (see Connection::disconnect).
+//!
+//! The counter is held via shared_ptr by the Connection and by every
+//! ProxyServer object created for the connection, because a ProxyServer
+//! object kept alive by an in-flight call can outlive the Connection (see
+//! ~ProxyServerBase), and its destructor must decrement state that is still
+//! valid.
+//!
+//! ProxyServer<Thread> and ProxyServer<ThreadMap> are separate
+//! specializations (not ProxyServerBase instances) and are intentionally not
+//! counted: every application method body runs on an interface ProxyServer,
+//! which is counted and stays alive for the duration of the body, so counting
+//! those is sufficient.
+struct ServerObjectTracker
+{
+    //! Called from the ProxyServerBase constructor (on the event loop thread).
+    void add()
+    {
+        const Lock lock(m_mutex);
+        m_count += 1;
+    }
+
+    //! Called from the ProxyServerBase destructor (on the event loop thread).
+    void remove()
+    {
+        {
+            const Lock lock(m_mutex);
+            assert(m_count > 0);
+            m_count -= 1;
+        }
+        m_cv.notify_all();
+    }
+
+    //! Return the current count. May be called from any thread.
+    size_t count() const
+    {
+        const Lock lock(m_mutex);
+        return m_count;
+    }
+
+    //! Block until no server objects remain.
+    void wait()
+    {
+        Lock lock(m_mutex);
+        m_cv.wait(lock.m_lock, [this]() MP_REQUIRES(m_mutex) { return m_count == 0; });
+    }
+
+    mutable Mutex m_mutex;
+    std::condition_variable m_cv;
+    size_t m_count MP_GUARDED_BY(m_mutex){0};
+};
+
 //! Object holding network & rpc state associated with either an incoming server
 //! connection, or an outgoing client connection. It must be created and destroyed
 //! on the event loop thread.
@@ -442,22 +512,55 @@ class Connection
 public:
     Connection(EventLoop& loop, kj::Own<kj::AsyncIoStream>&& stream_)
         : m_loop(loop), m_stream(kj::mv(stream_)),
-          m_network(*m_stream, ::capnp::rpc::twoparty::Side::CLIENT, ::capnp::ReaderOptions()),
-          m_rpc_system(::capnp::makeRpcClient(m_network)) {}
+          m_network(std::in_place, *m_stream, ::capnp::rpc::twoparty::Side::CLIENT, ::capnp::ReaderOptions()),
+          m_rpc_system(::capnp::makeRpcClient(*m_network)) {}
     Connection(EventLoop& loop,
         kj::Own<kj::AsyncIoStream>&& stream_,
         const std::function<::capnp::Capability::Client(Connection&)>& make_client)
         : m_loop(loop), m_stream(kj::mv(stream_)),
-          m_network(*m_stream, ::capnp::rpc::twoparty::Side::SERVER, ::capnp::ReaderOptions()),
-          m_rpc_system(::capnp::makeRpcServer(m_network, make_client(*this))) {}
-
-    //! Run cleanup functions. Must be called from the event loop thread. First
-    //! calls synchronous cleanup functions while blocked (to free capnp
-    //! Capability::Client handles owned by ProxyClient objects), then schedules
-    //! asynchronous cleanup functions to run in a worker thread (to run
-    //! destructors of m_impl instances owned by ProxyServer objects).
+          m_network(std::in_place, *m_stream, ::capnp::rpc::twoparty::Side::SERVER, ::capnp::ReaderOptions()),
+          m_rpc_system(::capnp::makeRpcServer(*m_network, make_client(*this))) {}
+
+    //! Destroy the connection. Calls disconnect() if it has not been called
+    //! already. Must be called from the event loop thread.
     ~Connection() noexcept(false);
 
+    //! Sever the connection without destroying this object: cancel any pending
+    //! onDisconnect handlers, cancel KJ promises for calls in progress, tear
+    //! down the RPC system (garbage collecting any server objects that are not
+    //! kept alive by in-flight calls), run synchronous cleanup functions
+    //! registered by client objects (releasing their capnp
+    //! Capability::Client handles), and release Thread capabilities so worker
+    //! threads are torn down. Safe to call more than once; the destructor
+    //! calls it automatically if it has not been called. Must be called from
+    //! the event loop thread.
+    //!
+    //! Note: disconnecting cancels the KJ promise of any call in progress, but
+    //! a C++ server method body that was already dispatched to a worker thread
+    //! (see ProxyServer<Thread>::post) is not interrupted by this and runs to
+    //! completion.
+    void disconnect();
+
+    //! Block until no ProxyServer objects associated with this connection
+    //! remain, i.e. until no server call body is still executing (see
+    //! ServerObjectTracker). Meant to be called after disconnect(): before it,
+    //! new server objects can still be created and idle server objects are
+    //! not garbage collected, so the count would not drain. Must NOT be called
+    //! from the event loop thread: in-flight call bodies need the event loop
+    //! to deliver their results before their server objects are destroyed, so
+    //! blocking the loop here would deadlock.
+    //!
+    //! This lets shutdown code ensure no IPC call body is still executing (and
+    //! dereferencing application state that is about to be freed) after
+    //! incoming connections are disconnected. See Ipc::disconnectIncoming and
+    //! https://github.com/bitcoin/bitcoin/issues/35845.
+    void waitDrained();
+
+    //! Number of live ProxyServer objects associated with this connection.
+    //! After disconnect(), a nonzero count means server call bodies are still
+    //! executing on worker threads. May be called from any thread.
+    size_t pendingServerObjects() const { return m_server_objects->count(); }
+
     //! Register synchronous cleanup function to run on event loop thread (with
     //! access to capnp thread local variables) when disconnect() is called.
     //! any new i/o.
@@ -473,7 +576,7 @@ class Connection
         // handler fires, do not call the function f right away, instead add it
         // to the EventLoop TaskSet to avoid "Promise callback destroyed itself"
         // error in the typical case where f deletes this Connection object.
-        m_on_disconnect.add(m_network.onDisconnect().then(
+        m_on_disconnect->add(m_network->onDisconnect().then(
             [f = std::forward<F>(f), this]() mutable { m_loop->m_task_set->add(kj::evalLater(kj::mv(f))); }));
     }
 
@@ -483,8 +586,24 @@ class Connection
     //! TaskSet used to cancel the m_network.onDisconnect() handler for remote
     //! disconnections, if the connection is closed locally first by deleting
     //! this Connection object.
-    kj::TaskSet m_on_disconnect{m_error_handler};
-    ::capnp::TwoPartyVatNetwork m_network;
+    std::optional<kj::TaskSet> m_on_disconnect{std::in_place, m_error_handler};
+    //! Wrapped in std::optional so disconnect() can destroy it (and m_stream
+    //! below) to sever the transport while this object stays alive. Closing
+    //! the stream is what makes the peer observe the disconnect: it reads EOF
+    //! and fails its outstanding calls with DISCONNECTED errors.
+    std::optional<::capnp::TwoPartyVatNetwork> m_network;
+
+    //! Tracker for live ProxyServer objects associated with this connection,
+    //! used by waitDrained(). Held via shared_ptr because ProxyServer objects
+    //! kept alive by in-flight calls can outlive the Connection (see
+    //! ServerObjectTracker and ~ProxyServerBase).
+    //!
+    //! Must be declared before m_rpc_system: constructing m_rpc_system runs
+    //! the make_client callback, which creates the bootstrap (Init) server
+    //! object, whose ProxyServerBase constructor registers itself with this
+    //! tracker.
+    std::shared_ptr<ServerObjectTracker> m_server_objects{std::make_shared<ServerObjectTracker>()};
+
     std::optional<::capnp::RpcSystem<::capnp::rpc::twoparty::VatId>> m_rpc_system;
 
     // ThreadMap interface client, used to create a remote server thread when an
@@ -511,6 +630,9 @@ class Connection
     //! will be empty if all ProxyClient are destroyed cleanly before the
     //! connection is destroyed.
     CleanupList m_sync_cleanup_fns;
+
+    //! Set once disconnect() has run. Only accessed on the event loop thread.
+    bool m_disconnected{false};
 };
 
 //! Vat id for server side of connection. Required argument to RpcSystem::bootStrap()
@@ -605,8 +727,14 @@ ProxyClientBase<Interface, Impl>::~ProxyClientBase() noexcept
 
 template <typename Interface, typename Impl>
 ProxyServerBase<Interface, Impl>::ProxyServerBase(std::shared_ptr<Impl> impl, Connection& connection)
-    : m_impl(std::move(impl)), m_context(&connection)
+    : m_impl(std::move(impl)), m_context(&connection), m_server_objects(connection.m_server_objects)
 {
+    // Register this object with the connection's live-object tracker. This
+    // runs on the event loop thread, so it is ordered before any connection
+    // teardown (which also runs on the event loop thread): code that
+    // disconnects the connection and then calls Connection::waitDrained() is
+    // guaranteed to see this object.
+    m_server_objects->add();
     MP_LOG(*m_context.loop, Log::Debug) << "Creating " << CxxTypeName(*this) << " " << this;
     assert(m_impl);
 }
@@ -654,6 +782,15 @@ ProxyServerBase<Interface, Impl>::~ProxyServerBase()
     }
     assert(m_context.cleanup_fns.empty());
     MP_LOG(*m_context.loop, Log::Debug) << "Destroying " << CxxTypeName(*this) << " " << this;
+    // Deregister this object from the connection's live-object tracker,
+    // through the shared m_server_objects handle since m_context.connection
+    // may be dangling here (see comment above). Done at the end of the
+    // destructor so a zero count means destruction fully completed. Note that
+    // any m_impl destruction scheduled through addAsyncCleanup above is NOT
+    // covered by the tracker: it runs later on the async cleanup thread, so
+    // Connection::waitDrained() waits for server call bodies, not for
+    // m_impl destructors.
+    m_server_objects->remove();
 }
 
 //! If the capnp interface defined a special "destroy" method, as described the
@@ -763,6 +900,19 @@ struct ThreadContext
     //! to assert false if there's an attempt to execute a blocking operation
     //! which could deadlock the thread.
     bool loop_thread = false;
+
+    //! Destructor which destroys the request_threads and callback_threads map
+    //! entries one at a time, removing each entry from its map while holding
+    //! Waiter::m_mutex, but destroying the removed ProxyClient<Thread> object
+    //! after releasing the mutex. Removing entries under the mutex is
+    //! necessary because event loop threads can concurrently remove map
+    //! entries when connections are broken (see SetThread cleanup function),
+    //! so the maps cannot be destroyed without locking as an implicit
+    //! destructor would do. Destroying ProxyClient<Thread> objects after
+    //! releasing the mutex is necessary to respect lock order and avoid
+    //! locking Waiter::m_mutex before EventLoop::m_mutex (see
+    //! "Synchronization note" above).
+    ~ThreadContext();
 };
 
 template<typename T, typename Fn>

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