Allow wrapped C++ methods to detect request cancellations

CMakeLists.txt

@@ -155,6 +155,7 @@ set(MP_PUBLIC_HEADERS
   include/mp/proxy-io.h
   include/mp/proxy-types.h
   include/mp/proxy.h
+  include/mp/type-cancel.h
   include/mp/type-char.h
   include/mp/type-chrono.h
   include/mp/type-context.h

doc/design.md

@@ -210,6 +210,37 @@ Subsequent requests will reuse the existing thread capabilities held in `callbac
    - Posts the work lambda to that thread's queue via `waiter->post(invoke)`
    - Cleans up the `request_threads` entry
 
+### Cancellation with Cancel
+
+Adding a `Cancel` parameter to a method that already has a `Context` parameter passes the C++ method a `CancelToken` argument that reports whether the request was abandoned, because the client disconnected, exited, or dropped the call's promise:
+
+```capnp
+waitForValue @25 (context :Proxy.Context, minValue :Int32, cancel :Proxy.Cancel) -> (result :Int32);
+```
+
+```cpp
+int FooImplementation::waitForValue(int min_value, CancelToken cancel)
+{
+    // Construct the callback before locking the mutex it uses.
+    const OnCancel wake{cancel, [this] {
+        const std::lock_guard<std::mutex> lock{m_wait_mutex};
+        m_wait_cv.notify_all();
+    }};
+    std::unique_lock<std::mutex> lock{m_wait_mutex};
+    m_wait_cv.wait(lock, [&] { return m_value >= min_value || cancel.canceled(); });
+    return m_value;
+}
+```
+
+Without this, a method blocked in a long wait keeps waiting after its client is gone, pinning its worker thread and delaying shutdown. libmultiprocess already discards results of canceled requests, so a method that observes `canceled()` can simply return.
+
+Notes:
+
+- The parameter carries no data. The client-side `CustomBuildField` sends an empty struct and ignores the client's `CancelToken` argument, which can be a default-constructed token (is never canceled). The server-side `PassField` constructs a token backed by the request's `CancelMonitor`.
+- `Cancel` must be declared after a `Context` parameter. Without a worker thread the method would run on the event loop thread, which is also the thread that delivers cancellations, so the token could never fire while the method runs. The code generator reports an error for this.
+- The server allows Cap'n Proto to cancel these calls when a client abandons them, equivalent to annotating the method with `$Cxx.allowCancellation`, but without requiring an annotation and while working on Cap'n Proto versions that predate it.
+- `OnCancel` callbacks run on the event loop thread and must not block. Construct an `OnCancel` before locking any mutex its callback locks, and do not destroy it from inside its own callback.
+
 ## Interface Definitions
 
 As explained in the [usage](usage.md) document, interface descriptions need to be consumed both by the _libmultiprocess_ code generator, and by C++ code that calls and implements the interfaces. The C++ code only needs to know about C++ arguments and return types, while the code generator only needs to know about capnp arguments and return types, but both need to know class and method names, so the corresponding `.h` and `.capnp` source files contain some of the same information, and have to be kept in sync manually when methods or parameters change. Despite the redundancy, reconciling the interface definitions is designed to be _straightforward_ and _safe_. _Straightforward_ because there is no need to write manual serialization code or use awkward intermediate types like [`UniValue`](https://github.com/bitcoin/bitcoin/blob/master/src/univalue/include/univalue.h) instead of native types. _Safe_ because if there are any inconsistencies between API and data definitions (even minor ones like using a narrow int data type for a wider int API input), there are errors at build time instead of errors or bugs at runtime.

include/mp/proxy-io.h

@@ -10,7 +10,9 @@
 
 #include <mp/proxy.capnp.h>
 
+#include <capnp/common.h>
 #include <capnp/rpc-twoparty.h>
+#include <capnp/schema.h>
 
 #include <assert.h>
 #include <algorithm>
@@ -23,6 +25,7 @@
 #include <memory>
 #include <optional>
 #include <sstream>
+#include <stdint.h>
 #include <string>
 #include <thread>
 
@@ -57,14 +60,16 @@ struct ServerInvokeContext : InvokeContext
     //! results structs if the request is canceled while the worker thread is
     //! reading params (`call_context.getParams()`) or writing results
     //! (`call_context.getResults()`).
-    Lock* cancel_lock{nullptr};
-    //! For IPC methods that execute asynchronously, not on the event-loop
-    //! thread, this is set to true if the IPC call was canceled by the client
-    //! or canceled by a disconnection. If the call runs on the event-loop
-    //! thread, it can't be canceled. This should be accessed with cancel_lock
-    //! held if it is not null, since in the asynchronous case it is accessed
-    //! from multiple threads.
-    bool request_canceled{false};
+    Lock* request_lock{nullptr};
+    //! The request's cancellation state, owned by its `CancelMonitor`. Null
+    //! for methods executing on the event-loop thread.
+    CancelState* cancel_state{nullptr};
+    //! Whether the IPC call was canceled by the client or by a disconnection.
+    //! Calls running on the event-loop thread cannot be canceled, so this
+    //! always returns false there. A false return is only reliable while
+    //! request_lock is held, which prevents the request's params and results
+    //! structs from being freed after the check.
+    bool request_canceled() const { return cancel_state && cancel_state->canceled(); }
 
     ServerInvokeContext(ProxyServer& proxy_server, CallContext& call_context, int req)
         : InvokeContext{*proxy_server.m_context.connection}, proxy_server{proxy_server}, call_context{call_context}, req{req}
@@ -632,6 +637,40 @@ ProxyServerBase<Interface, Impl>::ProxyServerBase(std::shared_ptr<Impl> impl, Co
     assert(m_impl);
 }
 
+template <typename Interface, typename Impl>
+auto ProxyServerBase<Interface, Impl>::dispatchCall(
+    uint64_t interface_id, uint16_t method_id, DispatchContext context) -> DispatchCallResult
+{
+    // A method allows cancellation if any of its parameters in the interface
+    // schema is an `mp.Cancel` struct. This lookup mirrors
+    // `capnp::DynamicCapability::Server::dispatchCall`.
+    bool allow_cancel{false};
+    KJ_IF_MAYBE(iface, capnp::Schema::from<Interface>().findSuperclass(interface_id)) {
+        auto methods{iface->getMethods()};
+        if (method_id < methods.size()) {
+            for (const auto field : methods[method_id].getParamType().getFields()) {
+                const auto type{field.getType()};
+                if (type.isStruct() && type.asStruct().getProto().getId() == capnp::typeId<Cancel>()) allow_cancel = true;
+            }
+        }
+    }
+
+    // Cap'n Proto 1.0 replaced the dynamic `CallContext::allowCancellation()`
+    // opt-in with a static flag on the dispatch result, so the code path
+    // depends on the Cap'n Proto version. See the "Breaking change" note in
+    // https://capnproto.org/news/2023-07-28-capnproto-1.0.html
+#if CAPNP_VERSION < 1000000
+    if (allow_cancel) context.allowCancellation();
+#endif
+
+    auto result{Interface::Server::dispatchCall(interface_id, method_id, kj::mv(context))};
+
+#if CAPNP_VERSION >= 1000000
+    if (allow_cancel) result.allowCancellation = true;
+#endif
+    return result;
+}
+
 //! ProxyServer destructor, called from the EventLoop thread by Cap'n Proto
 //! garbage collection code after there are no more references to this object.
 //! This will typically happen when the corresponding ProxyClient object on the

include/mp/proxy-types.h

@@ -521,22 +521,22 @@ struct ServerCall
     template <typename ServerContext, typename... Args>
     decltype(auto) invoke(ServerContext& server_context, TypeList<>, Args&&... args) const
     {
-        // If cancel_lock is set, release it while executing the method, and
+        // If request_lock is set, release it while executing the method, and
         // reacquire it afterwards. The lock is needed to prevent params and
         // response structs from being deleted by the event loop thread if the
         // request is canceled, so it is only needed before and after method
         // execution. It is important to release the lock during execution
         // because the method can take arbitrarily long to return and the event
         // loop will need the lock itself in on_cancel if the call is canceled.
-        if (server_context.cancel_lock) server_context.cancel_lock->m_lock.unlock();
+        if (server_context.request_lock) server_context.request_lock->m_lock.unlock();
         return TryFinally(
             [&]() -> decltype(auto) {
                 return ProxyServerMethodTraits<
                     typename decltype(server_context.call_context.getParams())::Reads
                 >::invoke(server_context, std::forward<Args>(args)...);
             },
             [&] {
-                if (server_context.cancel_lock) server_context.cancel_lock->m_lock.lock();
+                if (server_context.request_lock) server_context.request_lock->m_lock.lock();
                 // If the IPC request was canceled, throw InterruptException
                 // because there is no point continuing and trying to fill the
                 // call_context.getResults() struct. It's also important to stop
@@ -549,7 +549,7 @@ struct ServerCall
                 // Since the call has been canceled that exception can't be
                 // returned to the caller, so it needs to be discarded like
                 // other result values.
-                if (server_context.request_canceled) throw InterruptException{"canceled"};
+                if (server_context.request_canceled()) throw InterruptException{"canceled"};
             });
     }
 };

include/mp/proxy.capnp

@@ -67,3 +67,12 @@ struct Context $count(0) {
    # Handle of the client thread that is calling the current method, and that
    # any callbacks made by the server thread should be made on.
 }
+
+struct Cancel {
+   # Method parameter mapping to a C++ CancelToken argument, which
+   # lets the method detect when its request was abandoned, either because
+   # the client disconnected or because the call’s promise was dropped.
+   #
+   # NOTE: Must be declared after a `mp.Context` parameter, which provides
+   # the worker thread the method will run on.
+}

include/mp/proxy.h

@@ -12,14 +12,22 @@
 #include <list>
 #include <memory>
 #include <stddef.h>
+#include <stdint.h>
 #include <tuple>
 #include <type_traits>
 #include <utility>
 #include <variant> // IWYU pragma: keep
 
+namespace capnp {
+class AnyPointer;
+template <typename Params, typename Results>
+class CallContext;
+} // namespace capnp
+
 namespace mp {
 class Connection;
 class EventLoop;
+
 //! Mapping from capnp interface type to proxy client implementation (specializations are generated by
 //! proxy-codegen.cpp).
 template <typename Interface> struct ProxyClient; // IWYU pragma: export
@@ -157,6 +165,21 @@ struct ProxyServerBase : public virtual Interface_::Server
     void invokeDestroy();
     using Interface_::Server::thisCap;
 
+    // Type aliases used by the `dispatchCall` signature below.
+    using DispatchCallResult = typename Interface_::Server::DispatchCallResult;
+    using DispatchContext = ::capnp::CallContext<::capnp::AnyPointer, ::capnp::AnyPointer>;
+    /**
+     * Override of the capnp-generated `dispatchCall` method, which the RPC
+     * system uses to route incoming requests to the corresponding method.
+     *
+     * By default, capnp runs each call to completion even after the
+     * client abandons it, so this override also permits cancellation for
+     * methods that declare an `mp.Cancel` parameter.
+     *
+     * Every generated `ProxyServer` inherits this class, so one override
+     * covers all interfaces. */
+    DispatchCallResult dispatchCall(uint64_t interface_id, uint16_t method_id, DispatchContext context) override;
+
     /**
      * Implementation pointer that may or may not be owned and deleted when this
      * capnp server goes out of scope. It is owned for servers created to wrap

include/mp/type-cancel.h

@@ -0,0 +1,42 @@
+// Copyright (c) The Bitcoin Core developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or http://www.opensource.org/licenses/mit-license.php.
+
+#ifndef MP_PROXY_TYPE_CANCEL_H
+#define MP_PROXY_TYPE_CANCEL_H
+
+#include <mp/proxy-io.h>
+#include <mp/util.h>
+
+namespace mp {
+//! CustomBuildField override for mp.Cancel arguments. The field is left
+//! empty because the server constructs the token from request state. A live
+//! token is ignored, so passing one probably means the caller wrongly
+//! expects cancellation to be forwarded.
+template <typename Value, typename Output>
+    requires FieldTypeIs<Output, Cancel::Builder>
+void CustomBuildField(TypeList<CancelToken>,
+                      Priority<1>,
+                      InvokeContext& invoke_context,
+                      Value&& value,
+                      Output&& output)
+{
+    if (value.m_state) {
+        MP_LOG(*invoke_context.connection.m_loop, Log::Warning)
+            << "Ignoring live CancelToken argument. Cancellation is not forwarded to the server.";
+    }
+    output.init();
+}
+
+//! PassField override for mp.Cancel arguments. Ignores the field and
+//! passes the method a `CancelToken` that observes the request's cancellation
+//! state, owned by its `CancelMonitor` and published in
+//! `server_context.cancel_state` by the Context's PassField (see type-context.h).
+template <typename Accessor, typename ServerContext, typename Fn, typename... Args>
+void PassField(Priority<1>, TypeList<CancelToken>, ServerContext& server_context, Fn&& fn, Args&&... args)
+{
+    fn.invoke(server_context, std::forward<Args>(args)..., CancelToken{server_context.cancel_state});
+}
+} // namespace mp
+
+#endif // MP_PROXY_TYPE_CANCEL_H

include/mp/type-context.h

@@ -10,6 +10,8 @@
 
 #include <kj/string.h>
 
+#include <optional>
+
 namespace mp {
 template <typename Output>
     requires FieldTypeIs<Output, Context::Builder>
@@ -97,20 +99,20 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
         auto& request_threads = thread_context.request_threads;
         ConnThread request_thread;
         bool inserted{false};
-        Mutex cancel_mutex;
-        Lock cancel_lock{cancel_mutex};
-        server_context.cancel_lock = &cancel_lock;
+        Mutex request_mutex;
+        Lock request_lock{request_mutex};
+        server_context.request_lock = &request_lock;
+        server_context.cancel_state = &cancel_monitor.m_state;
+        // Cancellation callback, registered in the loop.sync call below
+        // after checking that the request was not canceled early.
+        std::optional<OnCancel> on_cancel;
         loop.sync([&] {
             // Detect request being canceled before it executes.
-            if (cancel_monitor.m_canceled) {
-                server_context.request_canceled = true;
-                return;
-            }
+            if (cancel_monitor.m_state.canceled()) return;
             // Detect request being canceled while it executes.
-            assert(!cancel_monitor.m_on_cancel);
-            cancel_monitor.m_on_cancel = [&loop, &server_context, &cancel_mutex, req]() {
+            on_cancel.emplace(CancelToken{&cancel_monitor.m_state}, [&loop, &request_mutex, req] {
                 MP_LOG(loop, Log::Info) << "IPC server request #" << req << " canceled while executing.";
-                // Lock cancel_mutex here to block the event loop
+                // Lock request_mutex here to block the event loop
                 // thread and prevent it from deleting the request's
                 // params and response structs while the execution
                 // thread is accessing them. Because this lock is
@@ -119,14 +121,13 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                 // protection from the event loop deleting the
                 // structs _before_ the execution thread acquires
                 // it. So in addition to locking the mutex, the
-                // execution thread always checks request_canceled
+                // execution thread always checks request_canceled()
                 // as well before accessing the structs.
-                Lock cancel_lock{cancel_mutex};
-                server_context.request_canceled = true;
-            };
+                Lock request_lock{request_mutex};
+            });
             // Update requests_threads map if not canceled. We know
             // the request is not canceled currently because
-            // cancel_monitor.m_canceled was checked above and this
+            // `cancel_monitor.m_state.canceled()` was checked above and this
             // code is running on the event loop thread.
             std::tie(request_thread, inserted) = SetThread(
                 GuardedRef{thread_context.waiter->m_mutex, request_threads}, server.m_context.connection,
@@ -140,29 +141,20 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
         // makes another IPC call), so avoid modifying the map.
         const bool erase_thread{inserted};
         KJ_DEFER(
-            // Release the cancel lock before calling loop->sync and
+            // Release the request lock before calling loop->sync and
             // waiting for the event loop thread, because if a
             // cancellation happened, it needs to run the on_cancel
-            // callback above. It's safe to release cancel_lock at
+            // callback above. It's safe to release request_lock at
             // this point because the fn.invoke() call below will be
             // finished and no longer accessing the params or
             // results structs.
-            cancel_lock.m_lock.unlock();
+            request_lock.m_lock.unlock();
             // Erase the request_threads entry on the event loop
             // thread with loop->sync(), so if the connection is
             // broken there is not a race between this thread and
             // the disconnect handler trying to destroy the thread
             // client object.
             loop.sync([&] {
-                // Clear cancellation callback. At this point the
-                // method invocation finished and the result is
-                // either being returned, or discarded if a
-                // cancellation happened. So we do not need to be
-                // notified of cancellations after this point. Also
-                // we do not want to be notified because
-                // cancel_mutex and server_context could be out of
-                // scope when it happens.
-                cancel_monitor.m_on_cancel = nullptr;
                 auto self_dispose{kj::mv(self)};
                 if (erase_thread) {
                     // Look up the thread again without using existing
@@ -179,7 +171,7 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                 }
             });
         );
-        if (server_context.request_canceled) {
+        if (server_context.request_canceled()) {
             MP_LOG(loop, Log::Info) << "IPC server request #" << req << " canceled before it could be executed";
         } else KJ_IF_MAYBE(exception, kj::runCatchingExceptions([&]{
             try {