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

include/mp/proxy-io.h

@@ -432,6 +432,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 +508,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 +572,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))); }));
     }
 
@@ -484,7 +583,23 @@ class Connection
     //! 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;
+    //! 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 +626,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 +723,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 +778,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

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

src/mp/proxy.cpp

@@ -113,6 +113,25 @@ Connection::~Connection() noexcept(false)
     // event loop thread, and if there was a remote disconnect, this is called
     // by an 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);
+    if (m_disconnected) return;
+    m_disconnected = true;
+
+    // Cancel pending onDisconnect handlers first. Severing the connection
+    // below completes m_network.onDisconnect() promises, and the registered
+    // handlers (see _Serve and ConnectStream) destroy this Connection object.
+    // That is redundant when disconnect() is called from the destructor, and
+    // harmful when disconnect() is called separately by code that keeps using
+    // the object afterwards (e.g. code waiting for in-flight calls to finish
+    // before destroying it).
+    m_on_disconnect.clear();
 
     // Try to cancel any calls that may be executing.
     m_canceler.cancel("Interrupted by disconnect");
@@ -200,12 +219,47 @@ Connection::~Connection() noexcept(false)
     // 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 owned by this connection, so idle worker
+    // threads are stopped and joined now instead of when this object is
+    // destroyed. (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.) Previously this happened
+    // implicitly when the m_thread_pool and m_thread_map members were
+    // destroyed; it is done explicitly here so disconnect() has the same
+    // effect whether or not the object is destroyed right away.
+    m_thread_pool.clear();
+    m_thread_map = nullptr;
+
+    // Destroy the network and close the stream. Closing the stream is what
+    // makes the peer observe the disconnect: it reads EOF and fails its
+    // outstanding calls with DISCONNECTED errors. Previously this happened
+    // implicitly when the m_network and m_stream members were destroyed; it
+    // must be done explicitly here because when disconnect() is called
+    // without destroying this object, nothing else severs the transport (the
+    // m_rpc_system.reset() call above stops reading from the stream but does
+    // not reliably close it), and the peer would not learn about the
+    // disconnect. The network is destroyed first since it references the
+    // stream.
+    m_network.reset();
+    m_stream = nullptr;
+}
+
+void Connection::waitDrained()
+{
+    // Blocking the event loop thread here would deadlock: in-flight call
+    // bodies sync() back to the event loop to deliver their results, and
+    // server objects are destroyed on the event loop thread.
+    assert(std::this_thread::get_id() != m_loop->m_thread_id);
+    m_server_objects->wait();
 }
 
 CleanupIt Connection::addSyncCleanup(std::function<void()> fn)

test/mp/test/test.cpp

@@ -425,6 +425,77 @@ KJ_TEST("Calling async IPC method, with server disconnect after cleanup")
     EXPECT_EXCEPTION(foo->callFnAsync(), "IPC client method call interrupted by disconnect.");
 }
 
+KJ_TEST("Waiting for in-flight server call to finish after disconnect")
+{
+    // Regression test for bitcoin/bitcoin#35845. Disconnecting a connection
+    // cancels the KJ promise of an in-flight call, but a C++ server method
+    // body already dispatched to a worker thread runs to completion. Verify
+    // that Connection::waitDrained() blocks until such a body finishes and its
+    // server object is destroyed, so shutdown code can wait for a disconnected
+    // connection to become quiescent before freeing state the body accesses.
+
+    std::promise<void> body_started, release_body;
+    TestSetup setup;
+    ProxyClient<messages::FooInterface>* foo = setup.client.get();
+    foo->initThreadMap();
+
+    // A server call body that signals when it starts and then blocks until the
+    // test releases it, so the in-flight state can be observed
+    // deterministically.
+    setup.server->m_impl->m_fn = [&] {
+        body_started.set_value();
+        release_body.get_future().get();
+    };
+
+    // Grab the server Connection object on the event loop thread before
+    // disconnecting. It stays valid until server_disconnect() destroys it
+    // below.
+    Connection* connection{nullptr};
+    foo->m_context.loop->sync([&] { connection = setup.server->m_context.connection; });
+
+    // Invoke the async method on a separate thread so its body blocks there
+    // while this thread makes assertions. callFnAsync() takes an mp.Context,
+    // so its body runs on a worker thread via ProxyServer<Thread>::post().
+    std::thread call_thread([&] {
+        EXPECT_EXCEPTION(foo->callFnAsync(), "IPC client method call interrupted by disconnect.");
+    });
+    body_started.get_future().get();
+
+    // The FooInterface server object is the connection's only counted server
+    // object, and its call body is executing.
+    KJ_EXPECT(connection->pendingServerObjects() == 1);
+
+    // Disconnect. This cancels the call's promise (the client above sees the
+    // disconnect error), but the body is still blocked on the worker thread,
+    // so its server object must still be alive.
+    foo->m_context.loop->sync([&] { connection->disconnect(); });
+    KJ_EXPECT(connection->pendingServerObjects() == 1);
+
+    // A drain must block while the body runs and return only once it
+    // finishes, which is what Ipc::disconnectIncoming relies on during
+    // shutdown.
+    std::atomic<bool> drained{false};
+    std::thread drain_thread([&] {
+        connection->waitDrained();
+        drained = true;
+    });
+
+    // The body is still blocked, so waitDrained() must not have returned.
+    std::this_thread::sleep_for(std::chrono::milliseconds(20));
+    KJ_EXPECT(!drained);
+
+    // Let the body finish; the drain should now complete.
+    release_body.set_value();
+    drain_thread.join();
+    KJ_EXPECT(drained);
+    KJ_EXPECT(connection->pendingServerObjects() == 0);
+    call_thread.join();
+
+    // Destroy the drained connection. (~Connection notices disconnect() has
+    // already run and does not tear things down twice.)
+    setup.server_disconnect();
+}
+
 KJ_TEST("Destroying ProxyClient<> with destroy method after peer disconnect")
 {
     // Regression test for bitcoin-core/libmultiprocess#219 where