test: Add more comprehensive DisconnectTest

Add more comprehensive disconnect test to test disconnect handling at more
points during server IPC execution to catch more bugs and catch them more
reliably in a way that does not rely on thread execution order.

Author: ryanofsky

include/mp/proxy-io.h

@@ -12,6 +12,7 @@
 
 #include <capnp/rpc-twoparty.h>
 
+#include <any>
 #include <assert.h>
 #include <condition_variable>
 #include <functional>
@@ -270,6 +271,9 @@ class EventLoop
 
     //! External context pointer.
     void* m_context;
+
+    //! External callback for hooks / tests.
+    std::function<void(std::any)> m_signal;
 };
 
 //! Single element task queue used to handle recursive capnp calls. (If server

include/mp/type-context.h

@@ -9,6 +9,11 @@
 #include <mp/util.h>
 
 namespace mp {
+struct SignalCallStart{};
+struct SignalCallSetup{};
+struct SignalCallTeardown{};
+struct SignalCallEnd{};
+
 template <typename Output>
 void CustomBuildField(TypeList<>,
     Priority<1>,
@@ -70,6 +75,7 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                 Context::Reader context_arg = Accessor::get(params);
                 ServerContext server_context{server, call_context, req};
                 {
+                    if (server.m_context.loop->m_signal) server.m_context.loop->m_signal(SignalCallStart{});
                     // Before invoking the function, store a reference to the
                     // callbackThread provided by the client in the
                     // thread_local.request_threads map. This way, if this
@@ -102,6 +108,7 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                     const bool erase_thread{inserted};
                     KJ_DEFER(if (erase_thread) {
                         std::unique_lock<std::mutex> lock(thread_context.waiter->m_mutex);
+                        if (server.m_context.loop->m_signal) server.m_context.loop->m_signal(SignalCallTeardown{});
                         // Call erase here with a Connection* argument instead
                         // of an iterator argument, because the `request_thread`
                         // iterator may be invalid if the connection is closed
@@ -111,8 +118,11 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                         // erases the thread from the map, and also because the
                         // ProxyServer<Thread> destructor calls
                         // request_threads.clear().
+                        if (server.m_context.loop->m_signal) server.m_context.loop->m_signal(SignalCallTeardown{});
                         request_threads.erase(server.m_context.connection);
-                    });
+                        }
+                        if (server.m_context.loop->m_signal) server.m_context.loop->m_signal(SignalCallEnd{});
+                    );
                     fn.invoke(server_context, args...);
                 }
                 KJ_IF_MAYBE(exception, kj::runCatchingExceptions([&]() {

src/mp/proxy.cpp

@@ -7,6 +7,7 @@
 #include <mp/proxy-io.h>
 #include <mp/proxy-types.h>
 #include <mp/proxy.capnp.h>
+#include <mp/type-context.h>
 #include <mp/type-threadmap.h>
 #include <mp/util.h>
 
@@ -308,6 +309,7 @@ bool EventLoop::done() const
 std::tuple<ConnThread, bool> SetThread(ConnThreads& threads, std::mutex& mutex, Connection* connection, const std::function<Thread::Client()>& make_thread)
 {
     const std::unique_lock<std::mutex> lock(mutex);
+    if (connection->m_loop->m_signal) connection->m_loop->m_signal(SignalCallSetup{});
     auto thread = threads.find(connection);
     if (thread != threads.end()) return {thread, false};
     thread = threads.emplace(

test/mp/test/foo-types.h

@@ -94,6 +94,41 @@ inline void CustomPassMessage(InvokeContext& invoke_context,
     fn(mut);
     builder.setMessage(mut.message + " return");
 }
+
+//! CustomBuildField for TestArg parameter. TestArg doesn't do anything special
+//! on the client side, so this does nothing. TestArg server-side behavior is
+//! implemented in CustomPassField below.
+template <typename Output>
+requires (std::is_same_v<decltype(std::declval<Output>().get()), test::messages::TestArg::Builder>)
+void CustomBuildField(TypeList<>,
+    Priority<1>,
+    ClientInvokeContext& invoke_context,
+    Output&& output)
+{
+}
+
+//! CustomPassField processing TestArg parameter by calling a start_hook()
+//! function which returns a bool promise, and continuing to execute the IPC if
+//! the promise value is true, aborting if it is false. It also calls an
+//! end_hook() function after the IPC call finishes, if it wasn't aborted.
+template <typename Accessor, typename ServerContext, typename Fn, typename... Args>
+requires (std::is_same_v<decltype(Accessor::get(std::declval<ServerContext>().call_context.getParams())), test::messages::TestArg::Reader>)
+auto CustomPassField(TypeList<>, ServerContext& server_context, const Fn& fn, Args&&... args)
+{
+    const auto& start_hook = server_context.proxy_server.m_impl->m_start_hook;
+    return start_hook().then([old=server_context, call_context=kj::mv(server_context.call_context), fn, args...](bool invoke_fn) mutable {
+        // If start hook returns false, skip calling IPC function.
+        if (!invoke_fn) return kj::Promise<typename ServerContext::CallContext>(kj::mv(call_context));
+        // If start hook returns true, continue calling IPC function and
+        // processing parameters/return values and calling end hook.
+        ServerContext server_context{old.proxy_server, call_context, old.req};
+        return fn.invoke(server_context, args...).then([old=server_context](ServerContext::CallContext call_context) {
+            ServerContext server_context{old.proxy_server, call_context, old.req};
+            const auto& end_hook = server_context.proxy_server.m_impl->m_end_hook;
+            return end_hook().then([call_context = kj::mv(call_context)] { return call_context; });
+        });
+    });
+}
 } // namespace mp
 
 #endif // MP_TEST_FOO_TYPES_H

test/mp/test/foo.capnp

@@ -30,7 +30,7 @@ interface FooInterface $Proxy.wrap("mp::test::FooImplementation") {
     passEnum @15 (arg :Int32) -> (result :Int32);
     passFn @16 (context :Proxy.Context, fn :FooFn) -> (result :Int32);
     callFn @17 () -> ();
-    callFnAsync @18 (context :Proxy.Context) -> ();
+    callFnAsync @18 (testArg :TestArg, context :Proxy.Context) -> ();
 }
 
 interface FooCallback $Proxy.wrap("mp::test::FooCallback") {
@@ -73,3 +73,7 @@ struct Pair(T1, T2) {
     first @0 :T1;
     second @1 :T2;
 }
+
+# Special argument used in disconnect test to implement ASYNC_START / ASYNC_END hooks.
+struct TestArg $Proxy.count(0) {
+}

test/mp/test/foo.h

@@ -82,6 +82,10 @@ class FooImplementation
     void callFn() { assert(m_fn); m_fn(); }
     void callFnAsync() { assert(m_fn); m_fn(); }
     std::function<void()> m_fn;
+    //! Hooks used by disconnect test ASYNC_START and ASYNC_END disconnects.
+    //! Former returns false to able IPC call without executing it.
+    std::function<kj::Promise<bool>()> m_start_hook;
+    std::function<kj::Promise<void>()> m_end_hook;
 };
 
 } // namespace test