proxy-io: Generalize ConnectStream / ServeStream for use in tests

include/mp/proxy-io.h

@@ -23,6 +23,7 @@
 #include <map>
 #include <memory>
 #include <optional>
+#include <ranges>
 #include <sstream>
 #include <string>
 #include <thread>
@@ -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<Connection>;
+
+    //! 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<kj::TaskSet> m_task_set;
 
     //! List of connections.
-    std::list<Connection> m_incoming_connections;
+    Connections m_incoming_connections;
 
     //! Logging options
     LogOptions m_log_opts;
@@ -864,41 +871,77 @@ kj::Promise<T> ProxyServer<Thread>::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 <typename InitInterface>
-std::unique_ptr<ProxyClient<InitInterface>> ConnectStream(EventLoop& loop, Stream stream)
+std::unique_ptr<ProxyClient<InitInterface>> ConnectStream(EventLoop& loop, Stream stream, bool destroy_connection = true)
 {
     typename InitInterface::Client init_client(nullptr);
     std::unique_ptr<Connection> connection;
     loop.sync([&] {
         connection = std::make_unique<Connection>(loop, kj::mv(stream));
         init_client = connection->m_rpc_system->bootstrap(ServerVatId().vat_id).castAs<InitInterface>();
     });
-    return std::make_unique<ProxyClient<InitInterface>>(
-        kj::mv(init_client), connection.release(), /* destroy_connection= */ true);
+    return std::make_unique<ProxyClient<InitInterface>>(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 <typename InitInterface, typename InitImpl, typename OnDisconnect>
-void _Serve(EventLoop& loop, kj::Own<kj::AsyncIoStream>&& stream, InitImpl& init, OnDisconnect&& on_disconnect)
+std::pair<ProxyServer<InitInterface>*, EventLoop::Connections::iterator> _Serve(EventLoop& loop,
+    kj::Own<kj::AsyncIoStream>&& stream,
+    std::shared_ptr<InitImpl> init,
+    OnDisconnect&& on_disconnect,
+    bool destroy_connection = true)
 {
+    ProxyServer<InitInterface>* 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<ProxyServer<InitInterface>>(std::shared_ptr<InitImpl>(&init, [](InitImpl*){}), connection);
+        auto proxy_server = kj::heap<ProxyServer<InitInterface>>(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<OnDisconnect>(on_disconnect)]() mutable {
+    it->onDisconnect([&loop, it, on_disconnect = std::forward<OnDisconnect>(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 <typename InitInterface, typename InitImpl, typename OnDisconnect>
+std::pair<ProxyServer<InitInterface>*, EventLoop::Connections::iterator> _Serve(EventLoop& loop,
+    kj::Own<kj::AsyncIoStream>&& stream,
+    InitImpl& init,
+    OnDisconnect&& on_disconnect,
+    bool destroy_connection = true)
+{
+    return _Serve<InitInterface>(loop, kj::mv(stream), std::shared_ptr<InitImpl>(&init, [](InitImpl*){}),
+        std::forward<OnDisconnect>(on_disconnect), destroy_connection);
 }
 
 struct Listener
@@ -936,11 +979,13 @@ void _Listen(const std::shared_ptr<Listener>& 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 <typename InitInterface, typename InitImpl>
-void ServeStream(EventLoop& loop, Stream stream, InitImpl& init)
+std::pair<ProxyServer<InitInterface>*, EventLoop::Connections::iterator> ServeStream(
+    EventLoop& loop, Stream stream, InitImpl&& init, bool destroy_connection = true)
 {
-    _Serve<InitInterface>(loop, kj::mv(stream), init, [] {});
+    return _Serve<InitInterface>(loop, kj::mv(stream), std::forward<InitImpl>(init), /*on_disconnect=*/ [] {}, destroy_connection);
 }
 
 //! Given listening socket identifier and an init object, handle incoming

test/mp/test/test.cpp

@@ -74,13 +74,18 @@ static_assert(std::is_integral_v<decltype(kMP_MINOR_VERSION)>, "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<void()> server_disconnect;
     std::function<void()> server_disconnect_later;
-    std::function<void()> server_on_disconnect;
     std::function<void()> client_disconnect;
     std::promise<std::unique_ptr<ProxyClient<messages::FooInterface>>> client_promise;
     std::unique_ptr<ProxyClient<messages::FooInterface>> 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<Connection>(loop, kj::mv(pipe.ends[0]), [&](Connection& connection) {
-                      auto server_proxy = kj::heap<ProxyServer<messages::FooInterface>>(
-                          std::make_shared<FooImplementation>(), connection);
-                      server = server_proxy;
-                      return capnp::Capability::Client(kj::mv(server_proxy));
-                  });
-              server_disconnect = [&] { loop.sync([&] { server_connection.reset(); }); };
+              auto server_result = ServeStream<messages::FooInterface>(
+                  loop, kj::mv(pipe.ends[0]), std::make_shared<FooImplementation>(), 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<Connection>(loop, kj::mv(pipe.ends[1]));
-              auto client_proxy = std::make_unique<ProxyClient<messages::FooInterface>>(
-                  client_connection->m_rpc_system->bootstrap(ServerVatId().vat_id).castAs<messages::FooInterface>(),
-                  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<messages::FooInterface>(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<messages::FooInterface>* 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<void> fn_called;