diff --git a/include/mp/proxy-io.h b/include/mp/proxy-io.h index 49b0611a..f3f3ee97 100644 --- a/include/mp/proxy-io.h +++ b/include/mp/proxy-io.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -299,6 +300,12 @@ class EventLoop //! Check if loop should exit. bool done() const MP_REQUIRES(m_mutex); + //! Type of m_incoming_connections list. + using Connections = std::list; + + //! View of incoming connections yielding Connection& for each entry. + auto incomingConnections() { 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; @@ -346,7 +353,7 @@ class EventLoop std::unique_ptr m_task_set; //! List of connections. - std::list m_incoming_connections; + Connections m_incoming_connections; //! Logging options LogOptions m_log_opts; @@ -864,8 +871,13 @@ kj::Promise ProxyServer::post(Fn&& fn) //! If the init interface declares a construct() method, creating the client //! calls it, so this function may block making an IPC call and may throw if //! the call fails. +//! +//! If destroy_connection is false, the returned client does not take +//! ownership of the connection, and the caller is responsible for +//! disconnecting and freeing it (accessible as client->m_context.connection) +//! whenever it is done with it. template -std::unique_ptr> ConnectStream(EventLoop& loop, Stream stream) +std::unique_ptr> ConnectStream(EventLoop& loop, Stream stream, bool destroy_connection = true) { typename InitInterface::Client init_client(nullptr); std::unique_ptr connection; @@ -873,32 +885,63 @@ std::unique_ptr> ConnectStream(EventLoop& loop, Strea connection = std::make_unique(loop, kj::mv(stream)); init_client = connection->m_rpc_system->bootstrap(ServerVatId().vat_id).castAs(); }); - return std::make_unique>( - kj::mv(init_client), connection.release(), /* destroy_connection= */ true); + return std::make_unique>(kj::mv(init_client), connection.release(), destroy_connection); } //! Given stream and init objects, construct a new ProxyServer object that //! handles requests from the stream by calling the init object. Embed the -//! ProxyServer in a Connection object that is stored and erased if -//! disconnected. This should be called from the event loop thread. +//! ProxyServer in a Connection object that is stored in +//! loop.m_incoming_connections. This should be called from the event loop +//! thread. Returns the new ProxyServer along with an iterator to its +//! Connection. +//! +//! If destroy_connection is false, the connection is not erased +//! automatically when the peer disconnects, and the caller is responsible +//! for erasing loop.m_incoming_connections at the returned iterator +//! whenever it is done with the connection. template -void _Serve(EventLoop& loop, kj::Own&& stream, InitImpl& init, OnDisconnect&& on_disconnect) +std::pair*, EventLoop::Connections::iterator> _Serve(EventLoop& loop, + kj::Own&& stream, + std::shared_ptr init, + OnDisconnect&& on_disconnect, + bool destroy_connection = true) { + ProxyServer* server = nullptr; loop.m_incoming_connections.emplace_front(loop, kj::mv(stream), [&](Connection& connection) { - // Disable deleter so proxy server object doesn't attempt to delete the - // init implementation when the proxy client is destroyed or - // disconnected. - return kj::heap>(std::shared_ptr(&init, [](InitImpl*){}), connection); + auto proxy_server = kj::heap>(std::move(init), connection); + server = proxy_server; + return capnp::Capability::Client(kj::mv(proxy_server)); }); auto it = loop.m_incoming_connections.begin(); MP_LOG(loop, Log::Info) << "IPC server: socket connected."; if (loop.testing_hook_connected) loop.testing_hook_connected(); - it->onDisconnect([&loop, it, on_disconnect = std::forward(on_disconnect)]() mutable { + it->onDisconnect([&loop, it, on_disconnect = std::forward(on_disconnect), destroy_connection]() mutable { MP_LOG(loop, Log::Info) << "IPC server: socket disconnected."; - loop.m_incoming_connections.erase(it); + if (destroy_connection) loop.m_incoming_connections.erase(it); on_disconnect(); if (loop.testing_hook_disconnected) loop.testing_hook_disconnected(); }); + return {server, it}; +} + +//! Overload of _Serve that takes a reference to the init object instead of a +//! shared_ptr. ProxyServer objects use shared_ptr's internally to optionally +//! take ownership of interfaces being served, and free them when clients are +//! no longer using them. Some libmultiprocess callers take advantage of this +//! and pass shared_ptr's directly. But other callers that do not give away +//! ownership are not required to use shared_ptr, so this overload converts +//! references they pass into shared_ptr's with empty deleters. This prevents +//! the ProxyServer object from deleting the init object when the client is +//! disconnected. +template +std::pair*, EventLoop::Connections::iterator> _Serve(EventLoop& loop, + kj::Own&& stream, + InitImpl& init, + OnDisconnect&& on_disconnect, + bool destroy_connection = true) +{ + return _Serve(loop, kj::mv(stream), std::shared_ptr(&init, [](InitImpl*){}), + std::forward(on_disconnect), destroy_connection); } struct Listener @@ -936,11 +979,13 @@ void _Listen(const std::shared_ptr& listener, EventLoop& loop, InitImp } //! Given a stream and an init object, handle requests on the stream by calling -//! methods on the Init object. +//! methods on the Init object. See _Serve for details on the return value and +//! the destroy_connection parameter. template -void ServeStream(EventLoop& loop, Stream stream, InitImpl& init) +std::pair*, EventLoop::Connections::iterator> ServeStream( + EventLoop& loop, Stream stream, InitImpl&& init, bool destroy_connection = true) { - _Serve(loop, kj::mv(stream), init, [] {}); + return _Serve(loop, kj::mv(stream), std::forward(init), /*on_disconnect=*/ [] {}, destroy_connection); } //! Given listening socket identifier and an init object, handle incoming diff --git a/test/mp/test/test.cpp b/test/mp/test/test.cpp index 5ecb7cc4..54b4dba5 100644 --- a/test/mp/test/test.cpp +++ b/test/mp/test/test.cpp @@ -74,13 +74,18 @@ static_assert(std::is_integral_v, "MP_MINOR_VERSION * this to be true to simplify shutdown and avoid needing to call * client_disconnect manually, but false allows testing more ProxyClient * behavior and the "IPC client method called after disconnect" code path. + * + * Also accepts an analogous server_owns_connection option controlling + * whether the server Connection is erased automatically when the peer + * disconnects. False is needed by tests that need to keep the server + * Connection alive across a disconnect notification, tearing it down + * explicitly later via server_disconnect() instead. */ class TestSetup { public: std::function server_disconnect; std::function server_disconnect_later; - std::function server_on_disconnect; std::function client_disconnect; std::promise>> client_promise; std::unique_ptr> client; @@ -89,7 +94,7 @@ class TestSetup //! not start until the other members are initialized. std::thread thread; - TestSetup(bool client_owns_connection = true) + TestSetup(bool client_owns_connection = true, bool server_owns_connection = true) : thread{[&] { EventLoop loop("mptest", [](mp::LogMessage log) { // Info logs are not printed by default, but will be shown with `mptest --verbose` @@ -98,36 +103,21 @@ class TestSetup }); auto pipe = loop.m_io_context.provider->newTwoWayPipe(); - auto server_connection = - std::make_unique(loop, kj::mv(pipe.ends[0]), [&](Connection& connection) { - auto server_proxy = kj::heap>( - std::make_shared(), connection); - server = server_proxy; - return capnp::Capability::Client(kj::mv(server_proxy)); - }); - server_disconnect = [&] { loop.sync([&] { server_connection.reset(); }); }; + auto server_result = ServeStream( + loop, kj::mv(pipe.ends[0]), std::make_shared(), server_owns_connection); + server = server_result.first; + EventLoop::Connections::iterator server_it = server_result.second; + server_disconnect = [&] { loop.sync([&] { loop.m_incoming_connections.erase(server_it); }); }; server_disconnect_later = [&] { assert(std::this_thread::get_id() == loop.m_thread_id); - loop.m_task_set->add(kj::evalLater([&] { server_connection.reset(); })); + loop.m_task_set->add(kj::evalLater([&] { loop.m_incoming_connections.erase(server_it); })); }; - // Set handler to destroy the server when the client disconnects. This - // is ignored if server_disconnect() is called instead. Tests can - // assign server_on_disconnect to override the default behavior of - // destroying the server connection as soon as the disconnect is - // detected (in which case they need to destroy it themselves, - // e.g. by calling server_disconnect(), so the event loop can - // exit). - server_on_disconnect = [&] { server_connection.reset(); }; - server_connection->onDisconnect([&] { server_on_disconnect(); }); - - auto client_connection = std::make_unique(loop, kj::mv(pipe.ends[1])); - auto client_proxy = std::make_unique>( - client_connection->m_rpc_system->bootstrap(ServerVatId().vat_id).castAs(), - client_connection.get(), /* destroy_connection= */ client_owns_connection); - if (client_owns_connection) { - (void)client_connection.release(); - } else { - client_disconnect = [&] { loop.sync([&] { client_connection.reset(); }); }; + + auto client_proxy = + ConnectStream(loop, kj::mv(pipe.ends[1]), client_owns_connection); + if (!client_owns_connection) { + Connection* client_connection = client_proxy->m_context.connection; + client_disconnect = [&] { loop.sync([&] { delete client_connection; }); }; } client_promise.set_value(std::move(client_proxy)); @@ -414,23 +404,20 @@ KJ_TEST("Calling async IPC method with a remote disconnect while results are bui // per-granule shadow history before a late reader comes along. // // The server Connection object is deliberately kept alive during all this - // by overriding server_on_disconnect: destroying it would cancel the - // in-flight request (Connection::~Connection calls m_canceler.cancel(), - // setting request_canceled) and the worker would throw InterruptException - // instead of proceeding into getResults(). Keeping it alive matches the - // window in the original report, where the worker races with capnp's own - // internal teardown, which runs before any onDisconnect notification. - - TestSetup setup{/*client_owns_connection=*/false}; + // by constructing TestSetup with server_owns_connection=false: destroying + // it would cancel the in-flight request (Connection::~Connection calls + // m_canceler.cancel(), setting request_canceled) and the worker would + // throw InterruptException instead of proceeding into getResults(). + // Keeping it alive matches the window in the original report, where the + // worker races with capnp's own internal teardown, which runs before any + // onDisconnect notification. The connection is + // destroyed explicitly at the end of the test instead. + + TestSetup setup{/*client_owns_connection=*/false, /*server_owns_connection=*/false}; ProxyClient* foo = setup.client.get(); KJ_EXPECT(foo->add(1, 2) == 3); foo->initThreadMap(); - // Keep the server Connection object alive when the disconnect is detected - // so the in-flight request is not canceled (see comment above). The - // connection is destroyed at the end of the test instead. - setup.server_on_disconnect = [] {}; - // Signaled by the worker thread when the method body runs, just before it // returns and the worker calls getResults() and serializes the results. std::promise fn_called;