proxy: add local connection limit to ListenConnections

Add an optional max_connections parameter to ListenConnections() so a listener can stop accepting new connections after reaching a local connection cap and resume accepting after an existing connection disconnects.

Implement the limit with listener-local state tracking the listening socket, maximum number of active connections, and whether an async accept() has already been posted. This keeps the limit scoped to the individual listener instead of introducing global EventLoop state.

Extend the dedicated listener test coverage to verify that with max_connections=1 the first client is accepted normally, a second client is not accepted while the first remains connected, and the second client is accepted after the first disconnects.

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

@@ -9,6 +9,9 @@ include.
 
 ## v12
 - 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.
 
 ## [v11.0](https://github.com/bitcoin-core/libmultiprocess/commits/v11.0)
 - Tolerates unexpected exceptions in event loop `post()` callbacks.

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>
@@ -834,8 +836,8 @@ 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>
-void _Serve(EventLoop& loop, kj::Own<kj::AsyncIoStream>&& stream, InitImpl& init)
+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
@@ -846,23 +848,69 @@ 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.";
     if (loop.testing_hook_connected) loop.testing_hook_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};
+};
+
 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);
+
+inline bool _ListenAtCapacity(const ListenState& state)
 {
-    auto* ptr = listener.get();
+    return state.max_connections && state.active_connections >= *state.max_connections;
+}
+
+//! Return an EventLoopRef if this is a capped listener, to keep the event loop
+//! alive while waiting for connections.
+//!
+//! For uncapped listeners, we do not hold an EventLoopRef because the event
+//! loop should automatically exit when the last active client disconnects (when
+//! m_num_clients drops to 0).
+//!
+//! For capped listeners, we stop calling ptr->accept() when at capacity. Since
+//! no accept task remains pending in the task set, we must hold an EventLoopRef
+//! while listening to keep the event loop alive until the limit is reached.
+inline std::unique_ptr<EventLoopRef> _MakeCappedListenerRef(EventLoop& loop, const ListenState& state)
+{
+    return state.max_connections && *state.max_connections > 0 ? std::make_unique<EventLoopRef>(loop) : nullptr;
+}
+
+//! Given init object and a state object containing a connection receiver, handle
+//! incoming connections by calling _Serve, to create ProxyServer objects and
+//! forward requests to the init object.
+template <typename InitInterface, typename InitImpl>
+void _Listen(EventLoop& loop, InitImpl& init, const std::shared_ptr<ListenState>& state)
+{
+    if (_ListenAtCapacity(*state)) return;
+
+    auto* ptr = state->listener.get();
+    // Bind the EventLoopRef lifetime to the pending accept task. If the listener
+    // reaches capacity, _Listen returns early, destroying the accept_ref and
+    // letting the event loop client count decrement.
+    auto accept_ref{_MakeCappedListenerRef(loop, *state)};
     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, accept_ref = std::move(accept_ref)](kj::Own<kj::AsyncIoStream>&& stream) mutable {
+            ++state->active_connections;
+            _Serve<InitInterface>(loop, kj::mv(stream), init, [&loop, &init, state] {
+                const bool resume_accept{_ListenAtCapacity(*state)};
+                assert(state->active_connections > 0);
+                --state->active_connections;
+                if (resume_accept) _Listen<InitInterface>(loop, init, state);
+            });
+            _Listen<InitInterface>(loop, init, state);
         }));
 }
 
@@ -872,18 +920,21 @@ template <typename InitInterface, typename InitImpl>
 void ServeStream(EventLoop& loop, int fd, InitImpl& init)
 {
     _Serve<InitInterface>(
-        loop, loop.m_io_context.lowLevelProvider->wrapSocketFd(fd, kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP), init);
+        loop,
+        loop.m_io_context.lowLevelProvider->wrapSocketFd(fd, kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP),
+        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));
     });
 }
 

test/mp/test/listen_tests.cpp

@@ -19,6 +19,7 @@
 #include <memory>
 #include <mp/proxy-io.h>
 #include <mutex>
+#include <optional>
 #include <ratio> // IWYU pragma: keep
 #include <stdexcept>
 #include <string>
@@ -119,19 +120,28 @@ class ClientSetup
 class ListenSetup
 {
 public:
-    ListenSetup()
-        : thread([this] {
+    explicit ListenSetup(std::optional<size_t> max_connections = std::nullopt)
+        : capped_listener(max_connections.has_value()), thread([this, max_connections] {
               EventLoop loop("mptest-server", [this](mp::LogMessage log) {
                   KJ_LOG(INFO, log.level, log.message);
                   if (log.level == mp::Log::Raise) throw std::runtime_error(log.message);
+                  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();
+                  }
               });
               loop.testing_hook_connected = [&] {
                   std::lock_guard<std::mutex> lock(counter_mutex);
                   ++connected_count;
                   counter_cv.notify_all();
               };
+              {
+                  std::lock_guard<std::mutex> lock(counter_mutex);
+                  event_loop = &loop;
+              }
               FooImplementation foo;
-              ListenConnections<messages::FooInterface>(loop, listener.release(), foo);
+              ListenConnections<messages::FooInterface>(loop, listener.release(), foo, max_connections);
               ready_promise.set_value();
               loop.loop();
           })
@@ -141,9 +151,23 @@ class ListenSetup
 
     ~ListenSetup()
     {
+        if (capped_listener) {
+            EventLoop* loop;
+            {
+                std::lock_guard<std::mutex> lock(counter_mutex);
+                loop = event_loop;
+            }
+            if (loop) loop->sync([&] { loop->m_task_set.reset(); });
+        }
         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);
@@ -154,15 +178,27 @@ class ListenSetup
         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);
+        const bool matched = counter_cv.wait_until(lock, deadline, [&] {
+            return disconnected_count >= expected_count;
+        });
+        KJ_REQUIRE(matched);
+    }
+
     UnixListener listener;
     std::promise<void> ready_promise;
+    bool capped_listener{false};
     std::mutex counter_mutex;
     std::condition_variable counter_cv;
+    EventLoop* event_loop{nullptr};
     size_t connected_count{0};
+    size_t disconnected_count{0};
     //! Thread variable should be after other struct members so the thread does
     //! not start until the other members are initialized.
     std::thread thread;
-    
 };
 
 KJ_TEST("ListenConnections accepts incoming connections")
@@ -174,6 +210,50 @@ KJ_TEST("ListenConnections accepts incoming connections")
     KJ_EXPECT(client->client->add(1, 2) == 3);
 }
 
+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();
+    setup.WaitForDisconnectedCount(2);
+    std::this_thread::sleep_for(std::chrono::milliseconds(100));
+
+    auto client3 = std::make_unique<ClientSetup>(setup.listener.Connect());
+    setup.WaitForConnectedCount(3);
+    KJ_EXPECT(client3->client->add(3, 4) == 7);
+}
+
+KJ_TEST("ListenConnections keeps capped listeners alive before reaching the limit")
+{
+    ListenSetup setup(/*max_connections=*/2);
+
+    auto client1 = std::make_unique<ClientSetup>(setup.listener.Connect());
+    setup.WaitForConnectedCount(1);
+    KJ_EXPECT(client1->client->add(1, 2) == 3);
+
+    client1.reset();
+    setup.WaitForDisconnectedCount(1);
+    std::this_thread::sleep_for(std::chrono::milliseconds(100));
+
+    auto client2 = std::make_unique<ClientSetup>(setup.listener.Connect());
+    setup.WaitForConnectedCount(2);
+    KJ_EXPECT(client2->client->add(2, 3) == 5);
+}
+
 } // namespace
 } // namespace test
 } // namespace mp