From 0ec4d0b77835c34bf6af4a765a6eaf017929aa7f Mon Sep 17 00:00:00 2001
From: Lasse Blaauwbroek <lasse@blaauwbroek.eu>
Date: Sat, 28 Oct 2023 07:03:42 +0200
Subject: [PATCH] Allow cancellation of all capability contexts

---
 capnp/helpers/capabilityHelper.h | 17 ++++++++++++++++-
 capnp/helpers/helpers.pxd        |  5 +++--
 capnp/lib/capnp.pyx              |  3 ++-
 test/test_capability.py          | 20 ++++++++++++++++++++
 4 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/capnp/helpers/capabilityHelper.h b/capnp/helpers/capabilityHelper.h
index fc98b70..4dcc9b4 100644
--- a/capnp/helpers/capabilityHelper.h
+++ b/capnp/helpers/capabilityHelper.h
@@ -75,10 +75,19 @@ class PythonInterfaceDynamicImpl final: public capnp::DynamicCapability::Server
   kj::Own<PyRefCounter> py_server;
   kj::Own<PyRefCounter> kj_loop;
 
+#if (CAPNP_VERSION_MAJOR < 1)
   PythonInterfaceDynamicImpl(capnp::InterfaceSchema & schema,
                              kj::Own<PyRefCounter> _py_server,
                              kj::Own<PyRefCounter> kj_loop)
-    : capnp::DynamicCapability::Server(schema), py_server(kj::mv(_py_server)), kj_loop(kj::mv(kj_loop)) { }
+    : capnp::DynamicCapability::Server(schema),
+      py_server(kj::mv(_py_server)), kj_loop(kj::mv(kj_loop)) { }
+#else
+  PythonInterfaceDynamicImpl(capnp::InterfaceSchema & schema,
+                             kj::Own<PyRefCounter> _py_server,
+                             kj::Own<PyRefCounter> kj_loop)
+    : capnp::DynamicCapability::Server(schema, { true }),
+      py_server(kj::mv(_py_server)), kj_loop(kj::mv(kj_loop)) { }
+#endif
 
   ~PythonInterfaceDynamicImpl() {
   }
@@ -87,6 +96,12 @@ class PythonInterfaceDynamicImpl final: public capnp::DynamicCapability::Server
                          capnp::CallContext< capnp::DynamicStruct, capnp::DynamicStruct> context);
 };
 
+inline void allowCancellation(capnp::CallContext<capnp::DynamicStruct, capnp::DynamicStruct> context) {
+#if (CAPNP_VERSION_MAJOR < 1)
+  context.allowCancellation();
+#endif
+}
+
 class PyAsyncIoStream: public kj::AsyncIoStream {
 public:
   kj::Own<PyRefCounter> protocol;
diff --git a/capnp/helpers/helpers.pxd b/capnp/helpers/helpers.pxd
index e405d2b..a6c15c6 100644
--- a/capnp/helpers/helpers.pxd
+++ b/capnp/helpers/helpers.pxd
@@ -1,7 +1,7 @@
 from capnp.includes.capnp_cpp cimport (
     Maybe, PyPromise, VoidPromise, RemotePromise,
-    DynamicCapability, InterfaceSchema, EnumSchema, StructSchema, DynamicValue, Capability, 
-    RpcSystem, MessageBuilder, Own, PyRefCounter, Node, DynamicStruct
+    DynamicCapability, InterfaceSchema, EnumSchema, StructSchema, DynamicValue, Capability,
+    RpcSystem, MessageBuilder, Own, PyRefCounter, Node, DynamicStruct, CallContext
 )
 
 from capnp.includes.schema_cpp cimport ByteArray
@@ -22,6 +22,7 @@ cdef extern from "capnp/helpers/capabilityHelper.h":
     PyPromise convert_to_pypromise(RemotePromise)
     PyPromise convert_to_pypromise(VoidPromise)
     VoidPromise taskToPromise(Own[PyRefCounter] coroutine, PyObject* callback)
+    void allowCancellation(CallContext context) except +reraise_kj_exception nogil
     void init_capnp_api()
 
 cdef extern from "capnp/helpers/rpcHelper.h":
diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx
index 298b9bd..8227d42 100644
--- a/capnp/lib/capnp.pyx
+++ b/capnp/lib/capnp.pyx
@@ -1927,7 +1927,7 @@ async def kj_loop():
 
 async def run(coro):
     """Ensure that the coroutine runs while the KJ event loop is running
-    
+
     This is a shortcut for wrapping the coroutine in a :py:meth:`capnp.kj_loop` context manager.
 
     :param coro: Coroutine to run
@@ -1949,6 +1949,7 @@ cdef class _CallContext:
     cdef CallContext * thisptr
 
     cdef _init(self, CallContext other):
+        helpers.allowCancellation(other)
         self.thisptr = new CallContext(move(other))
         return self
 
diff --git a/test/test_capability.py b/test/test_capability.py
index 27fae44..4dc5dc3 100644
--- a/test/test_capability.py
+++ b/test/test_capability.py
@@ -1,4 +1,5 @@
 import pytest
+import asyncio
 
 import capnp
 import test_capability_capnp as capability
@@ -389,3 +390,22 @@ async def test_generic():
     obj = capnp._MallocMessageBuilder().get_root_as_any()
     obj.set_as_text("anypointer_")
     assert (await client.foo(obj)).b == "anypointer_test"
+
+
+class CancelServer(capability.TestInterface.Server):
+    def __init__(self, val=1):
+        self.val = val
+
+    async def foo(self, i, j, **kwargs):
+        with pytest.raises(asyncio.CancelledError):
+            await asyncio.sleep(10)
+
+
+async def test_cancel2():
+    client = capability.TestInterface._new_client(CancelServer())
+
+    task = asyncio.ensure_future(client.foo(1, True))
+    await asyncio.sleep(0)  # Make sure that the task runs
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task