Commit 359722f3 authored by Callum Ryan's avatar Callum Ryan Committed by Facebook GitHub Bot

Task cancellation from Python

Summary:
Extending the current support for python futures <-> folly tasks interop to also support cancellation.

By adding a new version of `bridgeCoroTask` that also accepts a `folly::CancellationToken` we can give python code the ability to cancel the folly task.

Reviewed By: nanshu

Differential Revision: D31765430

fbshipit-source-id: 4277e94e296f56b998c9c583823a93ce5b9e7586
parent ea4721b1
......@@ -25,6 +25,7 @@
#include <folly/Executor.h>
#include <folly/Portability.h>
#include <folly/CancellationToken.h>
#include <folly/experimental/coro/Task.h>
#include <folly/python/AsyncioExecutor.h>
#include <folly/python/executor.h>
......@@ -39,7 +40,8 @@ void bridgeCoroTask(
folly::Executor* executor,
folly::coro::Task<T>&& coroFrom,
folly::Function<void(folly::Try<T>&&, PyObject*)> callback,
PyObject* userData) {
PyObject* userData,
folly::CancellationToken&& cancellationToken) {
// We are handing over a pointer to a python object to c++ and need
// to make sure it isn't removed by python in that time.
Py_INCREF(userData);
......@@ -48,7 +50,22 @@ void bridgeCoroTask(
[callback = std::move(callback), userData, guard = std::move(guard)](
folly::Try<T>&& result) mutable {
callback(std::move(result), userData);
});
},
std::move(cancellationToken));
}
template <typename T>
void bridgeCoroTask(
folly::Executor* executor,
folly::coro::Task<T>&& coroFrom,
folly::Function<void(folly::Try<T>&&, PyObject*)> callback,
PyObject* userData) {
bridgeCoroTask(
executor,
std::move(coroFrom),
std::move(callback),
userData,
folly::CancellationToken());
}
template <typename T>
......
......@@ -14,11 +14,21 @@
from cpython.ref cimport PyObject
from folly cimport cFollyExecutor, cFollyTry
from libcpp cimport bool
cdef extern from "folly/python/coro.h" namespace "folly::coro" nogil:
cdef cppclass cFollyCoroTask "folly::coro::Task"[T]:
pass
cdef extern from "folly/CancellationToken.h" namespace "folly" nogil:
cdef cppclass cFollyCancellationToken "folly::CancellationToken":
pass
cdef cppclass cFollyCancellationSource "folly::CancellationSource":
cFollyCancellationSource() except +
cFollyCancellationToken getToken()
bool requestCancellation()
cdef extern from "folly/python/coro.h" namespace "folly::python":
void bridgeCoroTask[T](
cFollyCoroTask[T]&& fut,
......@@ -32,3 +42,10 @@ cdef extern from "folly/python/coro.h" namespace "folly::python":
void(*)(cFollyTry[T]&&, PyObject*),
PyObject* pyFuture
)
void bridgeCoroTaskWithCancellation "folly::python::bridgeCoroTask"[T](
cFollyExecutor* executor,
cFollyCoroTask[T]&& fut,
void(*)(cFollyTry[T]&&, PyObject*),
PyObject* pyFuture,
cFollyCancellationToken&& cancelToken,
)
......@@ -25,3 +25,15 @@ class Futures(unittest.TestCase):
loop = asyncio.get_event_loop()
res = loop.run_until_complete(simplebridgecoro.get_value_x5_coro(val))
self.assertEqual(val * 5, res)
def test_cancellation(self):
loop = asyncio.get_event_loop()
res = loop.run_until_complete(self._return_five_after_cancelled())
self.assertEqual(5, res)
async def _return_five_after_cancelled(self):
task = asyncio.create_task(simplebridgecoro.return_five_after_cancelled())
await asyncio.sleep(0.1)
self.assertFalse(task.done())
task.cancel()
return await task
......@@ -13,14 +13,16 @@
# limitations under the License.
import asyncio
from folly.coro cimport cFollyCoroTask, bridgeCoroTask
from folly.coro cimport cFollyCoroTask, bridgeCoroTask, bridgeCoroTaskWithCancellation, cFollyCancellationSource
from folly cimport cFollyTry
from folly.executor cimport get_executor
from libc.stdint cimport uint64_t
from cpython.ref cimport PyObject
from cython.operator cimport dereference as deref
cdef extern from "folly/python/test/simplecoro.h" namespace "folly::python::test":
cdef cFollyCoroTask[uint64_t] coro_getValueX5(uint64_t val)
cdef cFollyCoroTask[uint64_t] coro_returnFiveAfterCancelled()
def get_value_x5_coro(int val):
......@@ -34,6 +36,24 @@ def get_value_x5_coro(int val):
return fut
async def return_five_after_cancelled():
cancellation_source = cFollyCancellationSource()
loop = asyncio.get_event_loop()
fut = loop.create_future()
bridgeCoroTaskWithCancellation[uint64_t](
get_executor(),
coro_returnFiveAfterCancelled(),
handle_uint64_t,
<PyObject *>fut,
cancellation_source.getToken(),
)
try:
return await asyncio.shield(fut)
except asyncio.CancelledError:
cancellation_source.requestCancellation()
return await fut
cdef void handle_uint64_t(cFollyTry[uint64_t]&& res, PyObject* userData):
future = <object> userData
if res.hasException():
......
......@@ -18,6 +18,8 @@
#include <cstdint>
#include <folly/CancellationToken.h>
#include <folly/experimental/coro/Baton.h>
#include <folly/experimental/coro/Task.h>
namespace folly {
......@@ -30,6 +32,15 @@ folly::coro::Task<uint64_t> coro_getValueX5(uint64_t val) {
}
co_return val * 5;
}
folly::coro::Task<uint64_t> coro_returnFiveAfterCancelled() {
folly::coro::Baton baton;
const folly::CancellationToken& ct =
co_await folly::coro::co_current_cancellation_token;
folly::CancellationCallback cb{ct, [&] { baton.post(); }};
co_await baton;
co_return 5;
}
} // namespace test
} // namespace python
} // namespace folly
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment