proxy-io.h: Add EventLoopRef RAII class handle addClient/removeClient refcounting

Author: ryanofsky

include/mp/proxy-io.h

@@ -190,7 +190,8 @@ class EventLoop
     //! other IPC calls.
     void startAsyncThread(std::unique_lock<std::mutex>& lock);
 
-    //! Add/remove remote client reference counts.
+    //! Add/remove remote client reference counts. Can use EventLoopRef
+    //! below to avoid calling these directly.
     void addClient(std::unique_lock<std::mutex>& lock);
     bool removeClient(std::unique_lock<std::mutex>& lock);
     //! Check if loop should exit.

include/mp/proxy.h

@@ -8,6 +8,7 @@
 #include <mp/util.h>
 
 #include <array>
+#include <cassert>
 #include <functional>
 #include <list>
 #include <stddef.h>
@@ -47,6 +48,26 @@ inline void CleanupRun(CleanupList& fns) {
     }
 }
 
+//! Event loop smart pointer automatically calling addClient and removeClient.
+//! If a lock pointer argument is passed, the specified lock will be used,
+//! otherwise EventLoop::m_mutex will be locked when needed.
+class EventLoopRef
+{
+public:
+    explicit EventLoopRef(EventLoop& loop, std::unique_lock<std::mutex>* lock = nullptr);
+    EventLoopRef(EventLoopRef&& other) noexcept : m_loop(other.m_loop) { other.m_loop = nullptr; }
+    EventLoopRef(const EventLoopRef&) = delete;
+    EventLoopRef& operator=(const EventLoopRef&) = delete;
+    EventLoopRef& operator=(EventLoopRef&&) = delete;
+    ~EventLoopRef() { reset(); }
+    EventLoop& operator*() const { assert(m_loop); return *m_loop; }
+    EventLoop* operator->() const { assert(m_loop); return m_loop; }
+    bool reset();
+
+    EventLoop* m_loop{nullptr};
+    std::unique_lock<std::mutex>* m_lock{nullptr};
+};
+
 //! Context data associated with proxy client and server classes.
 struct ProxyContext
 {

include/mp/util.h

@@ -18,6 +18,7 @@
 #include <tuple>
 #include <type_traits>
 #include <utility>
+#include <variant>
 #include <vector>
 
 namespace mp {
@@ -130,6 +131,20 @@ const char* TypeName()
     return short_name ? short_name + 1 : display_name;
 }
 
+//! Convenient wrapper around std::variant<T*, T>
+template <typename T>
+struct PtrOrValue {
+    std::variant<T*, T> data;
+
+    template <typename... Args>
+    PtrOrValue(T* ptr, Args&&... args) : data(ptr ? ptr : std::variant<T*, T>{std::in_place_type<T>, std::forward<Args>(args)...}) {}
+
+    T& operator*() { return data.index() ? std::get<T>(data) : *std::get<T*>(data); }
+    T* operator->() { return &**this; }
+    T& operator*() const { return data.index() ? std::get<T>(data) : *std::get<T*>(data); }
+    T* operator->() const { return &**this; }
+};
+
 //! Analog to std::lock_guard that unlocks instead of locks.
 template <typename Lock>
 struct UnlockGuard

src/mp/proxy.cpp

@@ -48,6 +48,23 @@ void LoggingErrorHandler::taskFailed(kj::Exception&& exception)
     m_loop.log() << "Uncaught exception in daemonized task.";
 }
 
+EventLoopRef::EventLoopRef(EventLoop& loop, std::unique_lock<std::mutex>* lock) : m_loop(&loop), m_lock(lock)
+{
+    auto loop_lock{PtrOrValue{m_lock, m_loop->m_mutex}};
+    m_loop->addClient(*loop_lock);
+}
+
+bool EventLoopRef::reset()
+{
+    bool done = false;
+    if (auto* loop{m_loop}) {
+        m_loop = nullptr;
+        auto loop_lock{PtrOrValue{m_lock, loop->m_mutex}};
+        done = loop->removeClient(*loop_lock);
+    }
+    return done;
+}
+
 Connection::~Connection()
 {
     // Shut down RPC system first, since this will garbage collect Server