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 @@ ...@@ -25,6 +25,7 @@
#include <folly/Executor.h> #include <folly/Executor.h>
#include <folly/Portability.h> #include <folly/Portability.h>
#include <folly/CancellationToken.h>
#include <folly/experimental/coro/Task.h> #include <folly/experimental/coro/Task.h>
#include <folly/python/AsyncioExecutor.h> #include <folly/python/AsyncioExecutor.h>
#include <folly/python/executor.h> #include <folly/python/executor.h>
...@@ -39,7 +40,8 @@ void bridgeCoroTask( ...@@ -39,7 +40,8 @@ void bridgeCoroTask(
folly::Executor* executor, folly::Executor* executor,
folly::coro::Task<T>&& coroFrom, folly::coro::Task<T>&& coroFrom,
folly::Function<void(folly::Try<T>&&, PyObject*)> callback, 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 // 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. // to make sure it isn't removed by python in that time.
Py_INCREF(userData); Py_INCREF(userData);
...@@ -48,7 +50,22 @@ void bridgeCoroTask( ...@@ -48,7 +50,22 @@ void bridgeCoroTask(
[callback = std::move(callback), userData, guard = std::move(guard)]( [callback = std::move(callback), userData, guard = std::move(guard)](
folly::Try<T>&& result) mutable { folly::Try<T>&& result) mutable {
callback(std::move(result), userData); 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> template <typename T>
......
...@@ -14,11 +14,21 @@ ...@@ -14,11 +14,21 @@
from cpython.ref cimport PyObject from cpython.ref cimport PyObject
from folly cimport cFollyExecutor, cFollyTry from folly cimport cFollyExecutor, cFollyTry
from libcpp cimport bool
cdef extern from "folly/python/coro.h" namespace "folly::coro" nogil: cdef extern from "folly/python/coro.h" namespace "folly::coro" nogil:
cdef cppclass cFollyCoroTask "folly::coro::Task"[T]: cdef cppclass cFollyCoroTask "folly::coro::Task"[T]:
pass 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": cdef extern from "folly/python/coro.h" namespace "folly::python":
void bridgeCoroTask[T]( void bridgeCoroTask[T](
cFollyCoroTask[T]&& fut, cFollyCoroTask[T]&& fut,
...@@ -32,3 +42,10 @@ cdef extern from "folly/python/coro.h" namespace "folly::python": ...@@ -32,3 +42,10 @@ cdef extern from "folly/python/coro.h" namespace "folly::python":
void(*)(cFollyTry[T]&&, PyObject*), void(*)(cFollyTry[T]&&, PyObject*),
PyObject* pyFuture 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): ...@@ -25,3 +25,15 @@ class Futures(unittest.TestCase):
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
res = loop.run_until_complete(simplebridgecoro.get_value_x5_coro(val)) res = loop.run_until_complete(simplebridgecoro.get_value_x5_coro(val))
self.assertEqual(val * 5, res) 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 @@ ...@@ -13,14 +13,16 @@
# limitations under the License. # limitations under the License.
import asyncio import asyncio
from folly.coro cimport cFollyCoroTask, bridgeCoroTask from folly.coro cimport cFollyCoroTask, bridgeCoroTask, bridgeCoroTaskWithCancellation, cFollyCancellationSource
from folly cimport cFollyTry from folly cimport cFollyTry
from folly.executor cimport get_executor
from libc.stdint cimport uint64_t from libc.stdint cimport uint64_t
from cpython.ref cimport PyObject from cpython.ref cimport PyObject
from cython.operator cimport dereference as deref from cython.operator cimport dereference as deref
cdef extern from "folly/python/test/simplecoro.h" namespace "folly::python::test": 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_getValueX5(uint64_t val)
cdef cFollyCoroTask[uint64_t] coro_returnFiveAfterCancelled()
def get_value_x5_coro(int val): def get_value_x5_coro(int val):
...@@ -34,6 +36,24 @@ def get_value_x5_coro(int val): ...@@ -34,6 +36,24 @@ def get_value_x5_coro(int val):
return fut 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): cdef void handle_uint64_t(cFollyTry[uint64_t]&& res, PyObject* userData):
future = <object> userData future = <object> userData
if res.hasException(): if res.hasException():
......
...@@ -18,6 +18,8 @@ ...@@ -18,6 +18,8 @@
#include <cstdint> #include <cstdint>
#include <folly/CancellationToken.h>
#include <folly/experimental/coro/Baton.h>
#include <folly/experimental/coro/Task.h> #include <folly/experimental/coro/Task.h>
namespace folly { namespace folly {
...@@ -30,6 +32,15 @@ folly::coro::Task<uint64_t> coro_getValueX5(uint64_t val) { ...@@ -30,6 +32,15 @@ folly::coro::Task<uint64_t> coro_getValueX5(uint64_t val) {
} }
co_return val * 5; 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 test
} // namespace python } // namespace python
} // namespace folly } // 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