proxy: add local connection limit to ListenConnections

Author: enirox001

doc/usage.md

@@ -10,7 +10,7 @@ _libmultiprocess_ is a library and code generator that allows calling C++ class
 
 The `*.capnp` data definition files are consumed by the _libmultiprocess_ code generator and each `X.capnp` file generates `X.capnp.c++`, `X.capnp.h`, `X.capnp.proxy-client.c++`, `X.capnp.proxy-server.c++`, `X.capnp.proxy-types.c++`, `X.capnp.proxy-types.h`, and `X.capnp.proxy.h` output files. The generated files include `mp::ProxyClient<Interface>` and `mp::ProxyServer<Interface>` class specializations for all the interfaces in the `.capnp` files. These allow methods on C++ objects in one process to be called from other processes over IPC sockets.
 
-The `ProxyServer` objects help translate IPC requests from a socket to method calls on a local object. The `ProxyServer` objects are just used internally by the `mp::ServeStream(loop, socket, wrapped_object)` and `mp::ListenConnections(loop, socket, wrapped_object)` functions, and aren't exposed externally. The `ProxyClient` classes are exposed, and returned from the `mp::ConnectStream(loop, socket)` function and meant to be used directly. The classes implement methods described in `.capnp` definitions, and whenever any method is called, a request with the method arguments is sent over the associated IPC connection, and the corresponding `wrapped_object` method on the other end of the connection is called, with the `ProxyClient` method blocking until it returns and forwarding back any return value to the `ProxyClient` method caller.
+The `ProxyServer` objects help translate IPC requests from a socket to method calls on a local object. The `ProxyServer` objects are just used internally by the `mp::ServeStream(loop, socket, wrapped_object)` and `mp::ListenConnections(loop, socket, wrapped_object[, max_connections])` functions, and aren't exposed externally. The `ProxyClient` classes are exposed, and returned from the `mp::ConnectStream(loop, socket)` function and meant to be used directly. The classes implement methods described in `.capnp` definitions, and whenever any method is called, a request with the method arguments is sent over the associated IPC connection, and the corresponding `wrapped_object` method on the other end of the connection is called, with the `ProxyClient` method blocking until it returns and forwarding back any return value to the `ProxyClient` method caller.
 
 ## Example
 

doc/versions.md

@@ -7,8 +7,14 @@ Library versions are tracked with simple
 Versioning policy is described in the [version.h](../include/mp/version.h)
 include.
 
-## v10
+## v11
 - Current unstable version.
+- Adds an optional per-listener `max_connections` parameter to `ListenConnections()`
+  so servers can stop accepting new connections when a local connection cap is reached,
+  and resume accepting after existing connections disconnect.
+
+## [v10.0](https://github.com/bitcoin-core/libmultiprocess/commits/v10.0)
+- Prior unstable version before local listener connection-limit support.
 
 ## [v9.0](https://github.com/bitcoin-core/libmultiprocess/commits/v9.0)
 - Fixes race conditions where worker thread could be used after destruction, where getParams() could be called after request cancel, and where m_on_cancel could be called after request finishes.

include/mp/proxy-io.h

@@ -13,7 +13,9 @@
 #include <capnp/rpc-twoparty.h>
 
 #include <assert.h>
+#include <algorithm>
 #include <condition_variable>
+#include <cstdlib>
 #include <functional>
 #include <kj/function.h>
 #include <map>
@@ -820,8 +822,19 @@ std::unique_ptr<ProxyClient<InitInterface>> ConnectStream(EventLoop& loop, int f
 //! 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.
+template <typename InitInterface, typename InitImpl, typename OnDisconnect>
+void _Serve(EventLoop& loop, kj::Own<kj::AsyncIoStream>&& stream, InitImpl& init, OnDisconnect&& on_disconnect);
+
 template <typename InitInterface, typename InitImpl>
 void _Serve(EventLoop& loop, kj::Own<kj::AsyncIoStream>&& stream, InitImpl& init)
+{
+    _Serve<InitInterface>(loop, kj::mv(stream), init, [] {});
+}
+
+//! Internal helper supporting additional disconnect actions for listeners that
+//! need to resume accepting connections when capacity becomes available.
+template <typename InitInterface, typename InitImpl, typename OnDisconnect>
+void _Serve(EventLoop& loop, kj::Own<kj::AsyncIoStream>&& stream, InitImpl& init, OnDisconnect&& on_disconnect)
 {
     loop.m_incoming_connections.emplace_front(loop, kj::mv(stream), [&](Connection& connection) {
         // Disable deleter so proxy server object doesn't attempt to delete the
@@ -831,23 +844,45 @@ void _Serve(EventLoop& loop, kj::Own<kj::AsyncIoStream>&& stream, InitImpl& init
     });
     auto it = loop.m_incoming_connections.begin();
     MP_LOG(loop, Log::Info) << "IPC server: socket connected.";
-    it->onDisconnect([&loop, it] {
+    it->onDisconnect([&loop, it, on_disconnect = std::forward<OnDisconnect>(on_disconnect)]() mutable {
         MP_LOG(loop, Log::Info) << "IPC server: socket disconnected.";
         loop.m_incoming_connections.erase(it);
+        on_disconnect();
     });
 }
 
 //! Given connection receiver and an init object, handle incoming connections by
 //! calling _Serve, to create ProxyServer objects and forward requests to the
 //! init object.
+struct ListenState
+{
+    explicit ListenState(kj::Own<kj::ConnectionReceiver>&& listener_, std::optional<size_t> max_connections_)
+        : listener(kj::mv(listener_)), max_connections(max_connections_) {}
+
+    kj::Own<kj::ConnectionReceiver> listener;
+    std::optional<size_t> max_connections;
+    size_t active_connections{0};
+    bool accept_pending{false};
+};
+
 template <typename InitInterface, typename InitImpl>
-void _Listen(EventLoop& loop, kj::Own<kj::ConnectionReceiver>&& listener, InitImpl& init)
+void _Listen(EventLoop& loop, InitImpl& init, const std::shared_ptr<ListenState>& state)
 {
-    auto* ptr = listener.get();
+    if (state->accept_pending) return;
+    if (state->max_connections && state->active_connections >= *state->max_connections) return;
+
+    state->accept_pending = true;
+    auto* ptr = state->listener.get();
     loop.m_task_set->add(ptr->accept().then(
-        [&loop, &init, listener = kj::mv(listener)](kj::Own<kj::AsyncIoStream>&& stream) mutable {
-            _Serve<InitInterface>(loop, kj::mv(stream), init);
-            _Listen<InitInterface>(loop, kj::mv(listener), init);
+        [&loop, &init, state](kj::Own<kj::AsyncIoStream>&& stream) mutable {
+            state->accept_pending = false;
+            ++state->active_connections;
+            _Serve<InitInterface>(loop, kj::mv(stream), init, [&loop, &init, state] {
+                assert(state->active_connections > 0);
+                --state->active_connections;
+                _Listen<InitInterface>(loop, init, state);
+            });
+            _Listen<InitInterface>(loop, init, state);
         }));
 }
 
@@ -863,12 +898,12 @@ void ServeStream(EventLoop& loop, int fd, InitImpl& init)
 //! Given listening socket file descriptor and an init object, handle incoming
 //! connections and requests by calling methods on the Init object.
 template <typename InitInterface, typename InitImpl>
-void ListenConnections(EventLoop& loop, int fd, InitImpl& init)
+void ListenConnections(EventLoop& loop, int fd, InitImpl& init, std::optional<size_t> max_connections = std::nullopt)
 {
     loop.sync([&]() {
-        _Listen<InitInterface>(loop,
+        _Listen<InitInterface>(loop, init, std::make_shared<ListenState>(
             loop.m_io_context.lowLevelProvider->wrapListenSocketFd(fd, kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP),
-            init);
+            max_connections));
     });
 }
 

include/mp/version.h

@@ -24,7 +24,7 @@
 //! pointing at the prior merge commit. The /doc/versions.md file should also be
 //! updated, noting any significant or incompatible changes made since the
 //! previous version.
-#define MP_MAJOR_VERSION 10
+#define MP_MAJOR_VERSION 11
 
 //! Minor version number. Should be incremented in stable branches after
 //! backporting changes. The /doc/versions.md file should also be updated to

test/mp/test/test.cpp

@@ -12,6 +12,7 @@
 #include <chrono>
 #include <condition_variable>
 #include <cstdint>
+#include <cstdlib>
 #include <cstring>
 #include <functional>
 #include <future>
@@ -27,19 +28,174 @@
 #include <mp/proxy-io.h>
 #include <mp/util.h>
 #include <mp/version.h>
+#include <mutex>
 #include <optional>
 #include <set>
+#include <sys/socket.h>
+#include <sys/un.h>
 #include <stdexcept>
 #include <string>
 #include <string_view>
 #include <thread>
 #include <type_traits>
+#include <unistd.h>
 #include <utility>
 #include <vector>
 
 namespace mp {
 namespace test {
 
+namespace {
+
+class UnixListener
+{
+public:
+    UnixListener()
+    {
+        char dir_template[] = "/tmp/mptest-listener-XXXXXX";
+        char* dir = mkdtemp(dir_template);
+        KJ_REQUIRE(dir != nullptr);
+        m_dir = dir;
+        m_path = m_dir + "/socket";
+
+        m_fd = socket(AF_UNIX, SOCK_STREAM, 0);
+        KJ_REQUIRE(m_fd >= 0);
+
+        sockaddr_un addr{};
+        addr.sun_family = AF_UNIX;
+        KJ_REQUIRE(m_path.size() < sizeof(addr.sun_path));
+        std::strncpy(addr.sun_path, m_path.c_str(), sizeof(addr.sun_path) - 1);
+        KJ_REQUIRE(bind(m_fd, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) == 0);
+        KJ_REQUIRE(listen(m_fd, SOMAXCONN) == 0);
+    }
+
+    ~UnixListener()
+    {
+        if (m_fd >= 0) close(m_fd);
+        if (!m_path.empty()) unlink(m_path.c_str());
+        if (!m_dir.empty()) rmdir(m_dir.c_str());
+    }
+
+    int release()
+    {
+        int fd = m_fd;
+        m_fd = -1;
+        return fd;
+    }
+
+    int Connect() const
+    {
+        int fd = socket(AF_UNIX, SOCK_STREAM, 0);
+        KJ_REQUIRE(fd >= 0);
+
+        sockaddr_un addr{};
+        addr.sun_family = AF_UNIX;
+        KJ_REQUIRE(m_path.size() < sizeof(addr.sun_path));
+        std::strncpy(addr.sun_path, m_path.c_str(), sizeof(addr.sun_path) - 1);
+        KJ_REQUIRE(connect(fd, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) == 0);
+        return fd;
+    }
+
+private:
+    int m_fd{-1};
+    std::string m_dir;
+    std::string m_path;
+};
+
+class ClientSetup
+{
+public:
+    explicit ClientSetup(int fd)
+        : thread([this, fd] {
+              EventLoop loop("mptest-client", [](mp::LogMessage log) {
+                  if (log.level == mp::Log::Raise) throw std::runtime_error(log.message);
+              });
+              client_promise.set_value(ConnectStream<messages::FooInterface>(loop, fd));
+              loop.loop();
+          })
+    {
+        client = client_promise.get_future().get();
+    }
+
+    ~ClientSetup()
+    {
+        client.reset();
+        thread.join();
+    }
+
+    std::promise<std::unique_ptr<ProxyClient<messages::FooInterface>>> client_promise;
+    std::unique_ptr<ProxyClient<messages::FooInterface>> client;
+    std::thread thread;
+};
+
+class ListenSetup
+{
+public:
+    explicit ListenSetup(size_t max_connections)
+        : thread([this, max_connections] {
+              EventLoop loop("mptest-server", [this](mp::LogMessage log) {
+                  if (log.level == mp::Log::Raise) throw std::runtime_error(log.message);
+                  if (log.message.find("IPC server: socket connected.") != std::string::npos) {
+                      std::lock_guard<std::mutex> lock(counter_mutex);
+                      ++connected_count;
+                      counter_cv.notify_all();
+                  } else if (log.message.find("IPC server: socket disconnected.") != std::string::npos) {
+                      std::lock_guard<std::mutex> lock(counter_mutex);
+                      ++disconnected_count;
+                      counter_cv.notify_all();
+                  }
+              });
+              FooImplementation foo;
+              ListenConnections<messages::FooInterface>(loop, listener.release(), foo, max_connections);
+              ready_promise.set_value();
+              loop.loop();
+          })
+    {
+        ready_promise.get_future().get();
+    }
+
+    ~ListenSetup()
+    {
+        thread.join();
+    }
+
+    size_t ConnectedCount()
+    {
+        std::lock_guard<std::mutex> lock(counter_mutex);
+        return connected_count;
+    }
+
+    void WaitForConnectedCount(size_t expected_count)
+    {
+        std::unique_lock<std::mutex> lock(counter_mutex);
+        const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5);
+        bool matched = counter_cv.wait_until(lock, deadline, [&] {
+            return connected_count >= expected_count;
+        });
+        KJ_REQUIRE(matched);
+    }
+
+    void WaitForDisconnectedCount(size_t expected_count)
+    {
+        std::unique_lock<std::mutex> lock(counter_mutex);
+        const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5);
+        bool matched = counter_cv.wait_until(lock, deadline, [&] {
+            return disconnected_count >= expected_count;
+        });
+        KJ_REQUIRE(matched);
+    }
+
+    UnixListener listener;
+    std::promise<void> ready_promise;
+    std::thread thread;
+    std::mutex counter_mutex;
+    std::condition_variable counter_cv;
+    size_t connected_count{0};
+    size_t disconnected_count{0};
+};
+
+} // namespace
+
 /** Check version.h header values */
 constexpr auto kMP_MAJOR_VERSION{MP_MAJOR_VERSION};
 constexpr auto kMP_MINOR_VERSION{MP_MINOR_VERSION};
@@ -481,5 +637,25 @@ KJ_TEST("Make simultaneous IPC calls on single remote thread")
     KJ_EXPECT(expected == 400);
 }
 
+KJ_TEST("ListenConnections enforces a local connection limit")
+{
+    ListenSetup setup(/*max_connections=*/1);
+
+    auto client1 = std::make_unique<ClientSetup>(setup.listener.Connect());
+    setup.WaitForConnectedCount(1);
+    KJ_EXPECT(client1->client->add(1, 2) == 3);
+
+    auto client2 = std::make_unique<ClientSetup>(setup.listener.Connect());
+    std::this_thread::sleep_for(std::chrono::milliseconds(100));
+    KJ_EXPECT(setup.ConnectedCount() == 1);
+
+    client1.reset();
+    setup.WaitForDisconnectedCount(1);
+    setup.WaitForConnectedCount(2);
+
+    KJ_EXPECT(client2->client->add(2, 3) == 5);
+    client2.reset();
+}
+
 } // namespace test
 } // namespace mp