bug: fix mptest hang, ProxyClient<Thread> deadlock in disconnect handler

include/mp/proxy-io.h

@@ -275,16 +275,16 @@ struct Waiter
     template <typename Fn>
     void post(Fn&& fn)
     {
-        const std::unique_lock<std::mutex> lock(m_mutex);
+        const Lock lock(m_mutex);
         assert(!m_fn);
         m_fn = std::forward<Fn>(fn);
         m_cv.notify_all();
     }
 
     template <class Predicate>
-    void wait(std::unique_lock<std::mutex>& lock, Predicate pred)
+    void wait(Lock& lock, Predicate pred)
     {
-        m_cv.wait(lock, [&] {
+        m_cv.wait(lock.m_lock, [&]() MP_REQUIRES(m_mutex) {
             // Important for this to be "while (m_fn)", not "if (m_fn)" to avoid
             // a lost-wakeup bug. A new m_fn and m_cv notification might be sent
             // after the fn() call and before the lock.lock() call in this loop
@@ -307,9 +307,9 @@ struct Waiter
     //! mutexes than necessary. This mutex can be held at the same time as
     //! EventLoop::m_mutex as long as Waiter::mutex is locked first and
     //! EventLoop::m_mutex is locked second.
-    std::mutex m_mutex;
+    Mutex m_mutex;
     std::condition_variable m_cv;
-    std::optional<kj::Function<void()>> m_fn;
+    std::optional<kj::Function<void()>> m_fn MP_GUARDED_BY(m_mutex);
 };
 
 //! Object holding network & rpc state associated with either an incoming server
@@ -540,7 +540,7 @@ using ConnThread = ConnThreads::iterator;
 // Retrieve ProxyClient<Thread> object associated with this connection from a
 // map, or create a new one and insert it into the map. Return map iterator and
 // inserted bool.
-std::tuple<ConnThread, bool> SetThread(ConnThreads& threads, std::mutex& mutex, Connection* connection, const std::function<Thread::Client()>& make_thread);
+std::tuple<ConnThread, bool> SetThread(GuardedRef<ConnThreads> threads, Connection* connection, const std::function<Thread::Client()>& make_thread);
 
 struct ThreadContext
 {
@@ -556,7 +556,7 @@ struct ThreadContext
     //! `callbackThread` argument it passes in the request, used by the server
     //! in case it needs to make callbacks into the client that need to execute
     //! while the client is waiting. This will be set to a local thread object.
-    ConnThreads callback_threads;
+    ConnThreads callback_threads MP_GUARDED_BY(waiter->m_mutex);
 
     //! When client is making a request to a server, this is the `thread`
     //! argument it passes in the request, used to control which thread on
@@ -565,7 +565,7 @@ struct ThreadContext
     //! by makeThread. If a client call is being made from a thread currently
     //! handling a server request, this will be set to the `callbackThread`
     //! request thread argument passed in that request.
-    ConnThreads request_threads;
+    ConnThreads request_threads MP_GUARDED_BY(waiter->m_mutex);
 
     //! Whether this thread is a capnp event loop thread. Not really used except
     //! to assert false if there's an attempt to execute a blocking operation

include/mp/proxy-types.h

@@ -617,7 +617,7 @@ void clientInvoke(ProxyClient& proxy_client, const GetRequest& get_request, Fiel
     const char* disconnected = nullptr;
     proxy_client.m_context.loop->sync([&]() {
         if (!proxy_client.m_context.connection) {
-            const std::unique_lock<std::mutex> lock(thread_context.waiter->m_mutex);
+            const Lock lock(thread_context.waiter->m_mutex);
             done = true;
             disconnected = "IPC client method called after disconnect.";
             thread_context.waiter->m_cv.notify_all();
@@ -644,7 +644,7 @@ void clientInvoke(ProxyClient& proxy_client, const GetRequest& get_request, Fiel
                 } catch (...) {
                     exception = std::current_exception();
                 }
-                const std::unique_lock<std::mutex> lock(thread_context.waiter->m_mutex);
+                const Lock lock(thread_context.waiter->m_mutex);
                 done = true;
                 thread_context.waiter->m_cv.notify_all();
             },
@@ -656,13 +656,13 @@ void clientInvoke(ProxyClient& proxy_client, const GetRequest& get_request, Fiel
                     proxy_client.m_context.loop->logPlain()
                         << "{" << thread_context.thread_name << "} IPC client exception " << kj_exception;
                 }
-                const std::unique_lock<std::mutex> lock(thread_context.waiter->m_mutex);
+                const Lock lock(thread_context.waiter->m_mutex);
                 done = true;
                 thread_context.waiter->m_cv.notify_all();
             }));
     });
 
-    std::unique_lock<std::mutex> lock(thread_context.waiter->m_mutex);
+    Lock lock(thread_context.waiter->m_mutex);
     thread_context.waiter->wait(lock, [&done]() { return done; });
     if (exception) std::rethrow_exception(exception);
     if (!kj_exception.empty()) proxy_client.m_context.loop->raise() << kj_exception;

include/mp/type-context.h

@@ -25,7 +25,7 @@ void CustomBuildField(TypeList<>,
     // Also store the Thread::Client reference in the callback_threads map so
     // future calls over this connection can reuse it.
     auto [callback_thread, _]{SetThread(
-        thread_context.callback_threads, thread_context.waiter->m_mutex, &connection,
+        GuardedRef{thread_context.waiter->m_mutex, thread_context.callback_threads}, &connection,
         [&] { return connection.m_threads.add(kj::heap<ProxyServer<Thread>>(thread_context, std::thread{})); })};
 
     // Call remote ThreadMap.makeThread function so server will create a
@@ -43,7 +43,7 @@ void CustomBuildField(TypeList<>,
         return request.send().getResult(); // Nonblocking due to capnp request pipelining.
     }};
     auto [request_thread, _1]{SetThread(
-        thread_context.request_threads, thread_context.waiter->m_mutex,
+        GuardedRef{thread_context.waiter->m_mutex, thread_context.request_threads},
         &connection, make_request_thread)};
 
     auto context = output.init();
@@ -90,7 +90,7 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                     auto& thread_context = g_thread_context;
                     auto& request_threads = thread_context.request_threads;
                     auto [request_thread, inserted]{SetThread(
-                        request_threads, thread_context.waiter->m_mutex,
+                        GuardedRef{thread_context.waiter->m_mutex, request_threads},
                         server.m_context.connection,
                         [&] { return context_arg.getCallbackThread(); })};
 
@@ -101,7 +101,7 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                     // makes another IPC call), so avoid modifying the map.
                     const bool erase_thread{inserted};
                     KJ_DEFER(if (erase_thread) {
-                        std::unique_lock<std::mutex> lock(thread_context.waiter->m_mutex);
+                        Lock lock(thread_context.waiter->m_mutex);
                         // 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

include/mp/util.h

@@ -182,6 +182,17 @@ class MP_SCOPED_CAPABILITY Lock {
     std::unique_lock<std::mutex> m_lock;
 };
 
+template<typename T>
+struct GuardedRef
+{
+    Mutex& mutex;
+    T& ref MP_GUARDED_BY(mutex);
+};
+
+// CTAD for Clang 16: GuardedRef{mutex, x} -> GuardedRef<decltype(x)>
+template <class U>
+GuardedRef(Mutex&, U&) -> GuardedRef<U>;
+
 //! Analog to std::lock_guard that unlocks instead of locks.
 template <typename Lock>
 struct UnlockGuard

src/mp/proxy.cpp

@@ -25,7 +25,6 @@
 #include <kj/memory.h>
 #include <map>
 #include <memory>
-#include <mutex>
 #include <optional>
 #include <stdexcept>
 #include <string>
@@ -305,15 +304,15 @@ bool EventLoop::done() const
     return m_num_clients == 0 && m_async_fns->empty();
 }
 
-std::tuple<ConnThread, bool> SetThread(ConnThreads& threads, std::mutex& mutex, Connection* connection, const std::function<Thread::Client()>& make_thread)
+std::tuple<ConnThread, bool> SetThread(GuardedRef<ConnThreads> threads, Connection* connection, const std::function<Thread::Client()>& make_thread)
 {
-    const std::unique_lock<std::mutex> lock(mutex);
-    auto thread = threads.find(connection);
-    if (thread != threads.end()) return {thread, false};
-    thread = threads.emplace(
+    const Lock lock(threads.mutex);
+    auto thread = threads.ref.find(connection);
+    if (thread != threads.ref.end()) return {thread, false};
+    thread = threads.ref.emplace(
         std::piecewise_construct, std::forward_as_tuple(connection),
         std::forward_as_tuple(make_thread(), connection, /* destroy_connection= */ false)).first;
-    thread->second.setDisconnectCallback([&threads, &mutex, thread] {
+    thread->second.setDisconnectCallback([threads, thread] {
         // Note: it is safe to use the `thread` iterator in this cleanup
         // function, because the iterator would only be invalid if the map entry
         // was removed, and if the map entry is removed the ProxyClient<Thread>
@@ -323,9 +322,9 @@ std::tuple<ConnThread, bool> SetThread(ConnThreads& threads, std::mutex& mutex,
         // thread client m_disconnect_cb member so thread client destructor does not
         // try to unregister this callback after connection is destroyed.
         // Remove connection pointer about to be destroyed from the map
-        const std::unique_lock<std::mutex> lock(mutex);
+        const Lock lock(threads.mutex);
         thread->second.m_disconnect_cb.reset();
-        threads.erase(thread);
+        threads.ref.erase(thread);
     });
     return {thread, true};
 }
@@ -364,7 +363,7 @@ ProxyServer<Thread>::~ProxyServer()
     assert(m_thread_context.waiter.get());
     std::unique_ptr<Waiter> waiter;
     {
-        const std::unique_lock<std::mutex> lock(m_thread_context.waiter->m_mutex);
+        const Lock lock(m_thread_context.waiter->m_mutex);
         //! Reset thread context waiter pointer, as shutdown signal for done
         //! lambda passed as waiter->wait() argument in makeThread code below.
         waiter = std::move(m_thread_context.waiter);
@@ -398,7 +397,7 @@ kj::Promise<void> ProxyServer<ThreadMap>::makeThread(MakeThreadContext context)
         g_thread_context.thread_name = ThreadName(m_connection.m_loop->m_exe_name) + " (from " + from + ")";
         g_thread_context.waiter = std::make_unique<Waiter>();
         thread_context.set_value(&g_thread_context);
-        std::unique_lock<std::mutex> lock(g_thread_context.waiter->m_mutex);
+        Lock lock(g_thread_context.waiter->m_mutex);
         // Wait for shutdown signal from ProxyServer<Thread> destructor (signal
         // is just waiter getting set to null.)
         g_thread_context.waiter->wait(lock, [] { return !g_thread_context.waiter; });