proxy-io.h: add Waiter::m_mutex thread safety annotations

Author: ryanofsky

include/mp/proxy-io.h

@@ -273,16 +273,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, [&] {
             // 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
@@ -305,7 +305,7 @@ 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;
 };
@@ -544,7 +544,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(ConnThreads& threads, Mutex& mutex, Connection* connection, const std::function<Thread::Client()>& make_thread);
 
 //! The thread_local ThreadContext g_thread_context struct provides information
 //! about individual threads and a way of communicating between them. Because
@@ -609,7 +609,7 @@ struct ThreadContext
     //! request thread argument passed in that request.
     //!
     //! Synchronization note: \ref callback_threads note applies here as well.
-    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

@@ -118,7 +118,7 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                             // without violating lock order.
                             ConnThreads::node_type removed;
                             {
-                                std::unique_lock<std::mutex> lock(thread_context.waiter->m_mutex);
+                                Lock lock(thread_context.waiter->m_mutex);
                                 removed = request_threads.extract(server.m_context.connection);
                             }
                         });

src/mp/proxy.cpp

@@ -25,7 +25,6 @@
 #include <kj/memory.h>
 #include <map>
 #include <memory>
-#include <mutex>
 #include <optional>
 #include <stdexcept>
 #include <string>
@@ -308,13 +307,13 @@ 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(ConnThreads& threads, Mutex& mutex, Connection* connection, const std::function<Thread::Client()>& make_thread)
 {
     assert(std::this_thread::get_id() == connection->m_loop->m_thread_id);
     ConnThread thread;
     bool inserted;
     {
-        const std::unique_lock<std::mutex> lock(mutex);
+        const Lock lock(mutex);
         std::tie(thread, inserted) = threads.try_emplace(connection);
     }
     if (inserted) {
@@ -331,7 +330,7 @@ std::tuple<ConnThread, bool> SetThread(ConnThreads& threads, std::mutex& mutex,
             thread->second->m_disconnect_cb.reset();
 
             // Remove connection pointer about to be destroyed from the map
-            const std::unique_lock<std::mutex> lock(mutex);
+            const Lock lock(mutex);
             threads.erase(thread);
         });
     }
@@ -373,7 +372,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);
@@ -407,7 +406,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; });