Skip to content

Commit 0ec4d0b

Browse files
LasseBlaauwbroekhaata
authored andcommitted
Allow cancellation of all capability contexts
1 parent 49bda5c commit 0ec4d0b

File tree

4 files changed

+41
-4
lines changed

4 files changed

+41
-4
lines changed

capnp/helpers/capabilityHelper.h

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,10 +75,19 @@ class PythonInterfaceDynamicImpl final: public capnp::DynamicCapability::Server
7575
kj::Own<PyRefCounter> py_server;
7676
kj::Own<PyRefCounter> kj_loop;
7777

78+
#if (CAPNP_VERSION_MAJOR < 1)
7879
PythonInterfaceDynamicImpl(capnp::InterfaceSchema & schema,
7980
kj::Own<PyRefCounter> _py_server,
8081
kj::Own<PyRefCounter> kj_loop)
81-
: capnp::DynamicCapability::Server(schema), py_server(kj::mv(_py_server)), kj_loop(kj::mv(kj_loop)) { }
82+
: capnp::DynamicCapability::Server(schema),
83+
py_server(kj::mv(_py_server)), kj_loop(kj::mv(kj_loop)) { }
84+
#else
85+
PythonInterfaceDynamicImpl(capnp::InterfaceSchema & schema,
86+
kj::Own<PyRefCounter> _py_server,
87+
kj::Own<PyRefCounter> kj_loop)
88+
: capnp::DynamicCapability::Server(schema, { true }),
89+
py_server(kj::mv(_py_server)), kj_loop(kj::mv(kj_loop)) { }
90+
#endif
8291

8392
~PythonInterfaceDynamicImpl() {
8493
}
@@ -87,6 +96,12 @@ class PythonInterfaceDynamicImpl final: public capnp::DynamicCapability::Server
8796
capnp::CallContext< capnp::DynamicStruct, capnp::DynamicStruct> context);
8897
};
8998

99+
inline void allowCancellation(capnp::CallContext<capnp::DynamicStruct, capnp::DynamicStruct> context) {
100+
#if (CAPNP_VERSION_MAJOR < 1)
101+
context.allowCancellation();
102+
#endif
103+
}
104+
90105
class PyAsyncIoStream: public kj::AsyncIoStream {
91106
public:
92107
kj::Own<PyRefCounter> protocol;

capnp/helpers/helpers.pxd

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from capnp.includes.capnp_cpp cimport (
22
Maybe, PyPromise, VoidPromise, RemotePromise,
3-
DynamicCapability, InterfaceSchema, EnumSchema, StructSchema, DynamicValue, Capability,
4-
RpcSystem, MessageBuilder, Own, PyRefCounter, Node, DynamicStruct
3+
DynamicCapability, InterfaceSchema, EnumSchema, StructSchema, DynamicValue, Capability,
4+
RpcSystem, MessageBuilder, Own, PyRefCounter, Node, DynamicStruct, CallContext
55
)
66

77
from capnp.includes.schema_cpp cimport ByteArray
@@ -22,6 +22,7 @@ cdef extern from "capnp/helpers/capabilityHelper.h":
2222
PyPromise convert_to_pypromise(RemotePromise)
2323
PyPromise convert_to_pypromise(VoidPromise)
2424
VoidPromise taskToPromise(Own[PyRefCounter] coroutine, PyObject* callback)
25+
void allowCancellation(CallContext context) except +reraise_kj_exception nogil
2526
void init_capnp_api()
2627

2728
cdef extern from "capnp/helpers/rpcHelper.h":

capnp/lib/capnp.pyx

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1927,7 +1927,7 @@ async def kj_loop():
19271927

19281928
async def run(coro):
19291929
"""Ensure that the coroutine runs while the KJ event loop is running
1930-
1930+
19311931
This is a shortcut for wrapping the coroutine in a :py:meth:`capnp.kj_loop` context manager.
19321932
19331933
:param coro: Coroutine to run
@@ -1949,6 +1949,7 @@ cdef class _CallContext:
19491949
cdef CallContext * thisptr
19501950

19511951
cdef _init(self, CallContext other):
1952+
helpers.allowCancellation(other)
19521953
self.thisptr = new CallContext(move(other))
19531954
return self
19541955

test/test_capability.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import pytest
2+
import asyncio
23

34
import capnp
45
import test_capability_capnp as capability
@@ -389,3 +390,22 @@ async def test_generic():
389390
obj = capnp._MallocMessageBuilder().get_root_as_any()
390391
obj.set_as_text("anypointer_")
391392
assert (await client.foo(obj)).b == "anypointer_test"
393+
394+
395+
class CancelServer(capability.TestInterface.Server):
396+
def __init__(self, val=1):
397+
self.val = val
398+
399+
async def foo(self, i, j, **kwargs):
400+
with pytest.raises(asyncio.CancelledError):
401+
await asyncio.sleep(10)
402+
403+
404+
async def test_cancel2():
405+
client = capability.TestInterface._new_client(CancelServer())
406+
407+
task = asyncio.ensure_future(client.foo(1, True))
408+
await asyncio.sleep(0) # Make sure that the task runs
409+
task.cancel()
410+
with pytest.raises(asyncio.CancelledError):
411+
await task

0 commit comments

Comments
 (0)