Allow cancellation of all capability contexts

capnp/helpers/capabilityHelper.h

@@ -75,10 +75,19 @@ public:
   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 @@ public:
                          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;

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
+    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":

capnp/lib/capnp.pyx

@@ -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
 

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