proxy-io.h: Add Connection disconnect and waitDrained methods

[ryanofsky]

Add Connection class disconnect and waitDrained methods to provide more flexibility when forcibly disconnecting from remote clients or servers.

Without these methods, the only way to forcibly close IPC connections is to delete Connection objects. This works but is not ideal because once a Connection object is gone, it is difficult to track state still associated with the connection, particularly:

• ProxyServer objects that may still be alive because they are executing asynchronous requests made before the disconnect. Without a way to track these objects, there is no generic way to wait for requests to finish existing after disconnecting. So individual IPC interfaces like the Bitcoin mining interface would need to implement custom synchronization to avoid race conditions during shutdown. Followup PR ipc: make ipc::disconnectIncoming wait for in-progress calls to complete bitcoin/bitcoin#35932 builds on this PR, calling the new waitDrained method introduced here to avoid IPC mining crashes on Bitcoin core shutdown without needing to change the mining code. A unit test is added here simulating these mining crashes.

• ProxyClient objects that contain pointers to Connection objects. Currently ProxyClient object need to register cleanup handlers with Connection objects to deal with Connections being deleted, which consumes memory and complicates ProxyClient shutdown logic. After this change, a followup PR will drop the cleanup handlers so Connection objects no longer need to track lists of ProxyClient objects associated with them. This is implemented in proxy-io: Reference-count Connection objects #336.

[DrahtBot]

The following sections might be updated with supplementary metadata relevant to reviewers and maintainers.

Reviews

See the guideline and AI policy for information on the review process.

Type	Reviewers
Concept ACK	xyzconstant

If your review is incorrectly listed, please copy-paste <!--meta-tag:bot-skip--> into the comment that the bot should ignore.

Conflicts

No conflicts as of last run.

[xyzconstant]

Concept ACK

[enirox001]

CI seems upset?

/home/runner/work/libmultiprocess/libmultiprocess/src/mp/proxy.cpp should add these lines:
#include <capnp/rpc-twoparty.h>  // for TwoPartyVatNetwork

[ryanofsky]

Updated 39ed2ca -> a40189f (pr/keepconn.1 -> pr/keepconn.2, compare) fixing iwyu errors and olddeps failure due to incompatibility with old capnproto versions which lack kj:::TaskSet::clear method https://github.com/bitcoin-core/libmultiprocess/actions/runs/31189931644/job/92903876132?pr=335

Added 1 commits a40189f -> 11929f1 (pr/keepconn.2 -> pr/keepconn.3, compare) to fix pre-existing ~ThreadContext() bug exposed by combination of new test in this PR and the onDisconnect handler added in #298 commit bb47369 https://github.com/bitcoin-core/libmultiprocess/actions/runs/31662352370/job/94329590873?pr=335

Updated 11929f1 -> 901a090 (pr/keepconn.3 -> pr/keepconn.4, compare) to fix iwyu error https://github.com/bitcoin-core/libmultiprocess/actions/runs/31666872134/job/94343269972?pr=335

[enirox001]

Code Review 901a090

Separating connection teardown from destruction and providing a server-call drain functioanlity is a good addition. The overall approach makes sense. I intend to review this more

I think the commit messages and code documentation are a bit too verbose. The explanations are nice to have, but it overexplains quite often, which ultimately makes it a bit harder to understand. Would suggest some revisions to the commit messages and the documentation to increase clarity

In commit 39cc757: ipc: add Connection::disconnect() separating teardown from destruction

I also think this is not exactly a behavior-neutral change; the commit message itself says Two details are new: as we now explicitly cancel m_on_disconnect handlers before severing the connection, and explicitly release m_thread_pool and m_thread_map during disconnect() rather than relying on member destruction.

The m_on_disconnect change is especially not something I would call behavior-neutral, as now we have to proactively cancel because Connection remains alive after the transport is severed and is no longer a consequence of destruction teardown. So even though the externally observable behaviour might seem unchanged, the lifetime and cancellation behaviour has changed, and I think that distinction matters

So the text saying

“This is a behavior-neutral refactor: the same steps run in the same order on destruction.”

is a bit misleading i think?

Also, in commit a40189f, there does not seem to be a clear commit title and description here; they are together

Left a few more suggestions and nits below

[enirox001]

In commit 39cc757: ipc: add Connection::disconnect() separating teardown from destruction

disconnect() sets m_disconnected = true; later on we clean everything up, so I am unsure, but if there was a scenario where one of the cleanups threw, it would not complete the rest. This might not be a problem, but another call to disconnect() would be a no-op.

I do not think all the operations after this can cause this to throw and lead to this, but shutdownWrite() might if it throws an exception other than the ones mentioned.

A simple fix is to set the m_disconnected = true only after all teardown that must run has completed.

index 0aaa58a..8b9f458 100644
--- a/src/mp/proxy.cpp
+++ b/src/mp/proxy.cpp
@@ -124,7 +124,6 @@ void Connection::disconnect()
     // the event loop thread, like the destructor.
     assert(std::this_thread::get_id() == m_loop->m_thread_id);
     if (m_disconnected) return;
-    m_disconnected = true;

     // Cancel pending onDisconnect handlers first. Severing the connection
     // below completes m_network.onDisconnect() promises, and the registered
@@ -253,6 +252,8 @@ void Connection::disconnect()
     // stream.
     m_network.reset();
     m_stream = nullptr;
+
+    m_disconnected = true;
 }

 void Connection::waitDrained()

or a better solution that make sure the the cleanup happens even if shutdownWrite fails?

[ryanofsky]

re: #335 (comment)

I think I want to drop the m_disconnected variable and just treat m_network being nullopt the same as m_disconnected being true, which I think should be equivalent to your suggestions.

It's also true that cleanup functions throwing is not something that this library handles very well generally, and could handle better in many cases.

[enirox001]

In 39cc757: ipc: add Connection::disconnect() separating teardown from destruction

Before this PR, when a remote side disconnected, libmultiprocess had callbacks that would eventually remove the Connection. It does not necessarily call the remove operation immediately; it can schedule it into another task set. The new disconnect() wants different behaviour. such that it will disconnect and then call waitDrained later on. So it tries to reset the m_on_disconnect callbacks.

But if the callback has already progressed one step further before reset happens, this violates the goal of this new system.

In aa49a11 this is made to use a weak_ptr, but I wonder if we should move those changes to this pr instead? Or rather, a small cancellation guard could be added to this PR such that it keeps the existing changes focused while preventing the potential regression.

A minimal change adding a weak cancelation token that has moved into the event loop queue.

index 1f77b26..d817eb6 100644
--- a/include/mp/proxy-io.h
+++ b/include/mp/proxy-io.h
@@ -576,8 +576,18 @@ public:
         // handler fires, do not call the function f right away, instead add it
         // to the EventLoop TaskSet to avoid "Promise callback destroyed itself"
         // error in the typical case where f deletes this Connection object.
+        const std::weak_ptr<void> guard{m_on_disconnect_guard};
         m_on_disconnect->add(m_network->onDisconnect().then(
-            [f = std::forward<F>(f), this]() mutable { m_loop->m_task_set->add(kj::evalLater(kj::mv(f))); }));
+            [f = std::forward<F>(f), guard, this]() mutable {
+                m_loop->m_task_set->add(kj::evalLater(
+                    [f = kj::mv(f), guard]() mutable {
+                        // The connection-owned TaskSet may have already handed
+                        // this callback to the event-loop TaskSet by the time
+                        // disconnect() cancels it. Only run it if the
+                        // connection has not been disconnected in between.
+                        if (guard.lock()) f();
+                    }));
+            }));
     }

     EventLoopRef m_loop;
@@ -587,6 +597,10 @@ public:
     //! disconnections, if the connection is closed locally first by deleting
     //! this Connection object.
     std::optional<kj::TaskSet> m_on_disconnect{std::in_place, m_error_handler};
+    //! Lifetime token checked by onDisconnect handlers after they are handed
+    //! off to the EventLoop TaskSet. Reset by disconnect() so a handler already
+    //! queued there cannot run after local teardown.
+    std::shared_ptr<void> m_on_disconnect_guard{std::make_shared<char>()};
     //! Wrapped in std::optional so disconnect() can destroy it (and m_stream
     //! below) to sever the transport while this object stays alive. Closing
     //! the stream is what makes the peer observe the disconnect: it reads EOF
diff --git a/src/mp/proxy.cpp b/src/mp/proxy.cpp
index 0aaa58a..06063f8 100644
--- a/src/mp/proxy.cpp
+++ b/src/mp/proxy.cpp
@@ -133,6 +133,7 @@ void Connection::disconnect()
     // harmful when disconnect() is called separately by code that keeps using
     // the object afterwards (e.g. code waiting for in-flight calls to finish
     // before destroying it).
+    m_on_disconnect_guard.reset();
     m_on_disconnect.reset();

     // Try to cancel any calls that may be executing.

This change closes the gap where resetting m_on_disconnect was too late because the callback had already moved into the EventLoop task set

[ryanofsky]

re: #335 (comment)

Hmm, this is an interesting finding. But if there is a bug here, it seems like a pre-existing one, not something caused by this change or made worse by it.

You're saying if a remote disconnect happens first, and the m_network->onDisconnect() callback executes, but the kj::evalLater callback inside it does not execute yet, and if within that interval, the local process decided to delete the connection, then the connection could be deleted twice.

This does seem like it might be possible, and I'd want to look into it a little more and write a test. I'd still be inclined to save a fix for a different PR, and I believe as you pointed out #336 might fix this.

[enirox001]

In 39cc757: ipc: add Connection::disconnect() separating teardown from destruction

The listener now keeps a counter of the active connections added in 39a10ce. When it is full, it stops accepting new connections, and when a client disconnects, a callback decreases the counter, and the listener can start accepting again.

But when the server calls disconnect() it cancels that callback. The connection closes, but the counter does not change, so the listener might think it is full and never accept another connection

Added this change to so that the listener count can be updated for every disconnect, while automatic deletion happens only for remote disconnects

index 1f77b26..30627ec 100644
--- a/include/mp/proxy-io.h
+++ b/include/mp/proxy-io.h
@@ -1016,10 +1016,12 @@ 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, on_disconnect = std::forward<OnDisconnect>(on_disconnect)]() mutable {
+    it->addSyncCleanup([on_disconnect = std::forward<OnDisconnect>(on_disconnect)]() mutable {
+        on_disconnect();
+    });
+    it->onDisconnect([&loop, it]() mutable {
         MP_LOG(loop, Log::Info) << "IPC server: socket disconnected.";
         loop.m_incoming_connections.erase(it);
-        on_disconnect();
         if (loop.testing_hook_disconnected) loop.testing_hook_disconnected();
     });
 }

This test could also be added to verify the above behaviour

index a9d4dca..240af3f 100644
--- a/test/mp/test/listen_tests.cpp
+++ b/test/mp/test/listen_tests.cpp
@@ -265,6 +265,29 @@ KJ_TEST("ListenConnections enforces a local connection limit")
     KJ_EXPECT(client3->client->add(3, 4) == 7);
 }

+KJ_TEST("ListenConnections resumes after a local disconnect")
+{
+    ListenSetup server(/*max_connections=*/1);
+
+    auto client1 = std::make_unique<ClientSetup>(server.listener.MakeConnectedSocket());
+    server.WaitForConnectedCount(1);
+    KJ_EXPECT(client1->client->add(1, 2) == 3);
+
+    auto client2 = std::make_unique<ClientSetup>(server.listener.MakeConnectedSocket());
+    (**server.m_loop_ref).sync([] {});
+    KJ_EXPECT(server.ConnectedCount() == 1);
+
+    EventLoop& loop{**server.m_loop_ref};
+    loop.sync([&] {
+        KJ_REQUIRE(loop.m_incoming_connections.size() == 1);
+        loop.m_incoming_connections.front().disconnect();
+        loop.m_incoming_connections.pop_front();
+    });
+
+    server.WaitForConnectedCount(2);
+    KJ_EXPECT(client2->client->add(2, 3) == 5);
+}
+
 KJ_TEST("ListenConnections accepts multiple connections")
 {
     // With max-connections=2, two clients should be accepted and usable at the

[ryanofsky]

re: #335 (comment)

Good catch and nice test!

[enirox001]

In commit 39cc757: ipc: add Connection::disconnect() separating teardown from destruction

Connection can now remain alive after it has been disconnected, but the class does not define which methods are safe to call afterwards.

Previously, disconnection meant destroying the whole object, but now the connection object still exists. This is needed so callers can use methods such as waitDrained, but other methods still behave as if the connection is active.

Could some documentation, assertion, or runtime check be helpful for this?

[enirox001]

In 631d8d9: ipc: add Connection::waitDrained() to wait for in-flight server calls

The documentation for the waitDrained method in proxy-io.h says it is meant to be called after disconnect() call, but does nothing to enforce it, i think we can assert that the disconnect method has been called before calling waitDrained as such

index 0aaa58a..6a318b0 100644
--- a/src/mp/proxy.cpp
+++ b/src/mp/proxy.cpp
@@ -261,6 +261,7 @@ void Connection::waitDrained()
     // bodies sync() back to the event loop to deliver their results, and
     // server objects are destroyed on the event loop thread.
     assert(std::this_thread::get_id() != m_loop->m_thread_id);
+    assert(m_disconnected);
     m_server_objects->wait();
 }

[enirox001]

In 631d8d9: ipc: add Connection::waitDrained() to wait for in-flight server calls

nit:
The name waitDrained could be clearer if it is called waitServerCallsDrained? or at least document its exact scope a bit more clearly

[enirox001]

In commit 092d1db: test: cover draining in-flight server call after disconnect

I think using an exact count of one is a bit brittle; the test only needs to establish that something remains in flight. This would be less implementation-specific

index 5bccb86..da47fde 100644
--- a/test/mp/test/test.cpp
+++ b/test/mp/test/test.cpp
@@ -463,13 +463,13 @@ KJ_TEST("Waiting for in-flight server call to finish after disconnect")

     // The FooInterface server object is the connection's only counted server
     // object, and its call body is executing.
-    KJ_EXPECT(connection->pendingServerObjects() == 1);
+    KJ_EXPECT(connection->pendingServerObjects() > 0);

     // Disconnect. This cancels the call's promise (the client above sees the
     // disconnect error), but the body is still blocked on the worker thread,
     // so its server object must still be alive.
     foo->m_context.loop->sync([&] { connection->disconnect(); });
-    KJ_EXPECT(connection->pendingServerObjects() == 1);
+    KJ_EXPECT(connection->pendingServerObjects() > 0);

     // A drain must block while the body runs and return only once it
     // finishes, which is what Ipc::disconnectIncoming relies on during

[enirox001]

In commit 092d1db: test: cover draining in-flight server call after disconnect

The 20 ms check is a bit too scheduler dependent, if the drain thread has not been scheduled during that interval, drained remains false even if waitDrained() is broken and would return immediately causing false pass.

We could add a hook here that runs only when ServerObject::wait() sees a non zero count and is about to wait

index 1f77b26..ab506b6 100644
--- a/include/mp/proxy-io.h
+++ b/include/mp/proxy-io.h
@@ -494,12 +494,14 @@ struct ServerObjectTracker
     void wait()
     {
         Lock lock(m_mutex);
+        if (m_count != 0 && testing_hook_wait) testing_hook_wait();
         m_cv.wait(lock.m_lock, [this]() MP_REQUIRES(m_mutex) { return m_count == 0; });
     }

     mutable Mutex m_mutex;
     std::condition_variable m_cv;
     size_t m_count MP_GUARDED_BY(m_mutex){0};
+    std::function<void()> testing_hook_wait;
 };

 //! Object holding network & rpc state associated with either an incoming server
diff --git a/test/mp/test/test.cpp b/test/mp/test/test.cpp
index 5bccb86..eec52f6 100644
--- a/test/mp/test/test.cpp
+++ b/test/mp/test/test.cpp
@@ -474,13 +474,17 @@ KJ_TEST("Waiting for in-flight server call to finish after disconnect")
     // A drain must block while the body runs and return only once it
     // finishes, which is what Ipc::disconnectIncoming relies on during
     // shutdown.
+    std::promise<void> drain_waiting;
+    connection->m_server_objects->testing_hook_wait = [&] { drain_waiting.set_value(); };
     std::atomic<bool> drained{false};
     std::thread drain_thread([&] {
         connection->waitDrained();
         drained = true;
     });

-    // The body is still blocked, so waitDrained() must not have returned.
+    // Wait until waitDrained() has observed the live server object and is
+    // about to block, then verify it does not return while the body is blocked.
+    drain_waiting.get_future().get();
     std::this_thread::sleep_for(std::chrono::milliseconds(20));
     KJ_EXPECT(!drained);

The test will then wait for that hook before starting the 20ms check. This ensures the drain thread has entered wait and observed pending work. This substantially reduces the possibility of a false positive

[enirox001]

In commit 901a090: Fix thread map teardown race causing use-after-free on disconnect

The Waiter documentation says

//! 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.

But the new commit says

//! Waiter::m_mutex must not be held when EventLoop::m_mutex is
//! acquired

It also says releasing the waiter mutex avoids locking the Waiter mutex before the EventLoop mutex, these rules cannot both be correct.

I think an actual order should be identified and updated here

[ryanofsky]

Thanks for the review! Great catches and suggestions. Just left some quick feedback below to make sure I didn't miss anything

[ryanofsky]

re: #335 (comment)

I think I want to drop the m_disconnected variable and just treat m_network being nullopt the same as m_disconnected being true, which I think should be equivalent to your suggestions.

It's also true that cleanup functions throwing is not something that this library handles very well generally, and could handle better in many cases.

[ryanofsky]

re: #335 (comment)

Hmm, this is an interesting finding. But if there is a bug here, it seems like a pre-existing one, not something caused by this change or made worse by it.

You're saying if a remote disconnect happens first, and the m_network->onDisconnect() callback executes, but the kj::evalLater callback inside it does not execute yet, and if within that interval, the local process decided to delete the connection, then the connection could be deleted twice.

This does seem like it might be possible, and I'd want to look into it a little more and write a test. I'd still be inclined to save a fix for a different PR, and I believe as you pointed out #336 might fix this.

[ryanofsky]

re: #335 (comment)

Good catch and nice test!