def _threaded_post_alarm(self, task: Future, request: Callable[..., None], *args: Any) -> None: """ Callback for the alarm queue. Check result of the future task and print a log in case something went wrong :param task: future task returned by the thread pool executor """ extra_info = "{}, {}".format(request, args) if not task.done(): if task.running(): logger.warning(log_messages.THREADED_REQUEST_HAS_LONG_RUNTIME, self.timeout, extra_info) else: logger.warning(log_messages.THREADED_REQUEST_IS_STALE, extra_info) task.cancel() else: try: result = task.result() logger.trace("Task {} completed with result: {}", extra_info, result) except CancelledError as e: logger.error(log_messages.TASK_CANCELLED, task, extra_info, e, exc_info=True) except Exception as e: # pylint: disable=broad-except logger.error(log_messages.TASK_FAILED, task, extra_info, e, exc_info=True)
def wrapper(kill_switch): if portage.getpid() == parent_pid: # thread in main process def done_callback(result): result.cancelled() or result.exception() or result.result() kill_switch.set() def start_coroutine(future): result = asyncio.ensure_future(coroutine_func(), loop=parent_loop) pending[id(result)] = result result.add_done_callback(done_callback) future.set_result(result) future = Future() parent_loop.call_soon_threadsafe(start_coroutine, future) kill_switch.wait() if not future.done(): future.cancel() raise asyncio.CancelledError elif not future.result().done(): future.result().cancel() raise asyncio.CancelledError else: return future.result().result() # child process loop = global_event_loop() try: return loop.run_until_complete(coroutine_func()) finally: loop.close()
def test_first_return_value(self): """Test that the value from the first returned non-cancelled future is returned""" future1 = Future() future2 = Future() future3 = Future() evt = self._make_event([self._exec.make_expected_function_call(future1), self._exec.make_expected_function_call(future2), self._exec.make_expected_function_call(future3)]) evt.go() start = time.time() self.assertIsNone(evt.first_result(1), "No completed futures, yet something returned") duration = time.time() - start # This one might be a bit too dependent on where the test is running. let's see how it goes... self.assertAlmostEqual(1.0, duration, msg="Incorrect timeout duration", delta=0.5) future1.cancel() self.assertIsNone(evt.first_result(0), "One (ignored) cancelled future, yet something returned") future3.set_result("WOOT") self.assertEquals("WOOT", evt.first_result(0), "Future completed, yet value not returned")
def _ready_callback(self, fn, mapping, proxy: Future, future: Future): """ Internally handles completion of dependencies """ with self._lock: if not proxy in self._pending: return del self._pending[proxy] if future.cancelled(): proxy.cancel() if not proxy.set_running_or_notify_cancel(): return exception = future.exception() if exception is not None: proxy.set_exception(exception) return if mapping is None: dependencies = future.result() else: dependencies = mapping(*future.result()) internal = self._executor.submit(fn, *dependencies) internal.add_done_callback(partial(self._done_callback, proxy))
async def _call_func(self, func: Callable, args: tuple, kwargs: Dict[str, Any], future: Future) -> None: def callback(f: Future) -> None: if f.cancelled(): self.call(scope.cancel) try: retval = func(*args, **kwargs) if iscoroutine(retval): with CancelScope() as scope: if future.cancelled(): scope.cancel() else: future.add_done_callback(callback) retval = await retval except self._cancelled_exc_class: future.cancel() except BaseException as exc: if not future.cancelled(): future.set_exception(exc) # Let base exceptions fall through if not isinstance(exc, Exception): raise else: if not future.cancelled(): future.set_result(retval) finally: scope = None # type: ignore[assignment]
def cancel_future(self, *, future: Future, exception: Exception) -> None: # cancel the task for this event if config.current.parallel_processing != "process": LOG.warning( f"Cancel task on boundary event was requested, " f"but the ADHESIVE_PARALLEL_PROCESSING is not set " f"to 'process', but '{config.current.parallel_processing}'. " f"The result of the task is ignored, but the thread " f"keeps running in the background." ) if future in self.futures: future_mapping = self.futures[future] LOG.warning(f"Cancelling active future '{future_mapping.description}'") # The futures are queried later since they are still in the `futures` map # inside the project executor. So to ensure they won't throw an exception, # that won't be able to find the owning events anymore - since they are # being removed now with the `DONE` call, we'll remove the futures from # polling. self.remove_future(future) else: LOG.warning(f"Cancelling active future '{future}'") future.cancel() # FIXME: hack. It seems that the `set_exception` is not needed, but # if cancel isn't called, the process isn't terminated. If it's called # after set_exception, again it isn't terminated. try: future.set_exception(exception) except Exception: pass
def test_job_done_callback_bails_out_if_canceled(self, mock_restore): future = Future() future.cancel() self.assertTrue(future.cancelled()) rsc = Mock() self.executor.job_done_callback(rsc, self.logger, future) self.assertEqual(mock_restore.call_args, call(rsc, self.logger)) self.assertTrue(self.executor.exceptions.empty())
def test_no_result_copy_cancel(self): source = Future() target = Future() copy(source, target, copy_result=False) source.cancel() self.assertTrue(target.cancel()) self.assertTrue(target.done())
def test_propagates_only_once_2(): rpc_fut = RPCFuture() fut = Future() rpc_fut.attach(fut) fut.set_result(42) fut.cancel() assert rpc_fut.result(timeout=1) == 42
def test_copy_cancel(self): source = Future() target = Future() copy(source, target) source.cancel() self.assertTrue(target.cancel()) self.assertTrue(target.done())
def test_no_cancel_copy_cancel(self): source = Future() target = Future() copy(source, target, copy_cancel=False) source.cancel() self.assertFalse(target.cancelled()) self.assertFalse(target.done())
async def _call_func(self, func: Callable, args: tuple, future: Future) -> None: try: retval = func(*args) if isinstance(retval, Coroutine): future.set_result(await retval) else: future.set_result(retval) except self._cancelled_exc_class: future.cancel() except BaseException as exc: future.set_exception(exc)
def _done_callback(self, proxy: Future, future: Future): """ Internally handles completion of executor future, copies result to proxy Args: fn (function): [description] future (Future): [description] """ if future.cancelled(): proxy.cancel() exception = future.exception() if exception is not None: proxy.set_exception(exception) else: result = future.result() proxy.set_result(result)
def test_await_all_futures(self): """Test that the await_all_futures call will return all complete, non-cancelled futures""" future1 = Future() future2 = Future() evt = self._make_event([self._exec.make_expected_function_call(future1), self._exec.make_expected_function_call(future2)]) evt.go() self.assertListEqual([], evt.await_all(0), "No ready futures, but something returned") future1.cancel() future2.set_result("WOOT") self.assertListEqual([future2], evt.await_all(0), "Single future was not returned after cancel/complete")
class Zipper(object): def __init__(self, fs): self.fs = list(fs) self.out = Future() self.done = False self.lock = Lock() self.count_remaining = len(self.fs) for (idx, future) in enumerate(self.fs): chain_cancel(self.out, future) future.add_done_callback( weak_callback(partial(self.handle_done, idx))) def handle_done(self, index, f): set_result = False set_exception = False cancel = False with self.lock: if self.done: pass elif f.cancelled(): self.done = True cancel = True elif f.exception(): self.done = True set_exception = True else: self.fs[index] = f.result() self.count_remaining -= 1 if self.count_remaining == 0: self.done = True set_result = True if cancel: self.out.cancel() if set_result: try_set_result(self.out, maketuple(self.fs)) if set_exception: copy_future_exception(f, self.out)
def _copy(self, src: Future, dst: Future) -> None: gotit = self._lock.acquire(blocking=False) if not gotit: return try: if self._fired: return self._fired = True if src.cancelled(): dst.cancel() return try: dst.set_result(src.result()) except Exception as e: dst.set_exception(e) finally: self._lock.release()
def test_zip_inner_cancel(): f_a = f_return("a") f_b = Future() f_c = Future() future = f_zip(f_a, f_b, f_c) assert f_b.cancel() # Cancelling the inner future should cause the outer # future to also become cancelled, as well as the # other inner futures. assert f_c.cancelled() assert future.cancelled()
class Section(ReprableMixin): """ A collection of orders. :param orders: a list of orders :param disposition: if Disposition.JOINABLE then this section can be joined with other sections. If Disposition.CANNOT_JOIN then all orders from this section will be executed before proceeding to next one """ _REPR_FIELDS = ('orders', 'disposition') __slots__ = ('orders', 'disposition', 'future') def __init__(self, orders: tp.List[Order] = None, disposition: Disposition = Disposition.JOINABLE): self.future = Future() self.orders = orders or [] self.disposition = disposition def cancel(self) -> bool: return self.future.cancel() def result(self, timeout: tp.Optional[float] = None): self.future.result(timeout) def __bool__(self) -> bool: return bool(self.orders) @classmethod def from_json(cls, dct: dict): return Section(orders_from_list(dct['orders']), Disposition(dct.get('disposition', 0))) def __iadd__(self, other: tp.Union[Order, 'Section']) -> 'Section': if isinstance(other, Order): self.orders.append(other) else: other.future = self.future self.orders.extend(other.orders) return self def is_joinable(self) -> bool: return self.disposition == Disposition.JOINABLE def max_wait(self) -> tp.Optional[float]: wait = None for order in (or_ for or_ in self.orders if isinstance(or_, WaitOrder)): if wait is None: wait = order.period else: if wait < order.period: wait = order.period return wait
class _Chain(object): """ A linked list of futures. Each future yields a result and the next link in the chain """ def __init__(self): self._next = Future() def push(self, value): """Sets the value of this link in the chain waking up all waiting listeners and returns a reference to the next link. """ next_ = _Chain() self._next.set_result((value, next_)) return next_ def close(self): """Finish the chain at this link. It will not given a value. All current and future listeners will be woken with a :exception:`ClosedError` and it will no longer be possible to add new links. """ self._next.cancel() def wait(self, timeout=None): try: result = self._next.result(timeout) except CancelledError: raise ClosedError() return result def wait_result(self, timeout=None): return self.wait(timeout)[0] def wait_next(self, timeout=None): return self.wait(timeout)[1]
class AsyncioExecutorFuture: def __init__(self, response): self._fut = Future() self._response = response response.add_done_callback(self._on_response_done) self._asyncio_task = None self._canceled = False def __getattr__(self, name): return getattr(self._fut, name) def cancel(self): if self._canceled: return self._canceled = True self._fut.cancel() return True def _on_task_done(self, asyncio_task): if self._canceled: return error = result = None try: error = asyncio_task.exception() if error is None: result = asyncio_task.result() except CancelledError: self._fut.cancel() else: if error is None: self._fut.set_result(result) else: self._fut.set_exception(error) def _on_response_done(self, response): if self._canceled: return try: asyncio_task = response.result() except CancelledError: self._fut.cancel() else: self._asyncio_task = asyncio_task asyncio_task.add_done_callback(self._on_task_done)
class ExpectationBase(ABC): always_monitor = False def __init__(self): self._future = Future() self._awaited = False self._scheduler = None self._success = False self._timeout = None self._deadline = None self._timedout = False # FIXME: float_tol should be moved to ArsdkExpectationBase self._float_tol = DEFAULT_FLOAT_TOL def _schedule(self, scheduler): # This expectation is scheduled on the `scheduler`, subclasses of ExpectationBase can # perform some operations on this scheduler: schedule another expectation later or # perform an operation on the scheduler object when this expectation is schedule (like # sending a message for which this expectation object expect some result). # IMPORTANT NOTE: this function (or its overridden versions) should be non-blocking self._awaited = True self._scheduler = scheduler if self._timeout is not None: self._deadline = timestamp_now() + self._timeout def success(self): return self._success def wait(self, _timeout=None): if self._awaited: try: self._future.result(timeout=_timeout) except FutureTimeoutError: self.set_timedout() except FutureCancelledError: self.cancel() return self def add_done_callback(self, cb): self._future.add_done_callback(lambda f: cb(self)) def set_success(self): if not self._future.done(): self._success = True self._future.set_result(self.received_events()) return True return False def set_exception(self, exception): if not self._future.done(): self._future.set_exception(exception) def set_timeout(self, _timeout): self._timeout = _timeout def set_timedout(self): if self._future.done(): return False if not self._success: self._timedout = True self.cancel() return True return False def cancel(self): if self._future.done(): return False self._future.cancel() return True def cancelled(self): return self._future.cancelled() def timedout(self): if self._timedout: return True if self._success: return False if self._deadline is not None: timedout = timestamp_now() > self._deadline if timedout: self.set_timedout() return self._timedout def set_float_tol(self, _float_tol): self._float_tol = _float_tol def base_copy(self, *args, **kwds): other = self.__class__(*args, **kwds) ExpectationBase.__init__(other) other._timeout = self._timeout other._float_tol = self._float_tol return other @abstractmethod def copy(self): """ All expectations sublclasses must implement a shallow copy. """ pass def done(self): return (self._future.done() or not self._awaited) and self._success def __bool__(self): return self.done() def __or__(self, other): return WhenAnyExpectation([self, other]) def __and__(self, other): return WhenAllExpectations([self, other]) def __rshift__(self, other): return WhenSequenceExpectations([self, other]) __nonzero__ = __bool__
class ExpectationBase(object): __metaclass__ = ABCMeta def __init__(self): self._future = Future() self._awaited = False self._scheduler = None self._success = False self._timeout = None self._deadline = None self._timedout = False self._float_tol = DEFAULT_FLOAT_TOL def _schedule(self, scheduler): # This expectation is scheduled on the `scheduler`, subclasses of ExpectationBase can # perform some operations on this scheduler: schedule another expectation later or # perform an operation on the scheduler object when this expectation is schedule (like # sending a message for which this expectation object expect some result). self._awaited = True self._scheduler = scheduler if self._timeout is not None: self._deadline = timestamp_now() + self._timeout def success(self): return self._success def wait(self, _timeout=None): if self._awaited: try: self._future.result(timeout=_timeout) except FutureTimeoutError: self.set_timedout() except FutureCancelledError: self.cancel() return self def set_result(self): self._success = True return self._future.set_result(self.received_events()) def set_exception(self, exception): return self._future.set_exception(exception) def set_timeout(self, _timeout): self._timeout = _timeout def set_timedout(self): if not self._success: self._timedout = True self.cancel() def cancel(self): return self._future.cancel() def cancelled(self): return self._future.cancelled() def timedout(self): if self._timedout: return True if self._success: return False if self._deadline is not None: self._timedout = (timestamp_now() > self._deadline) if self._timedout: self.cancel() return self._timedout def set_float_tol(self, _float_tol): self._float_tol = _float_tol def base_copy(self, *args, **kwds): other = self.__class__(*args, **kwds) ExpectationBase.__init__(other) other._timeout = self._timeout other._float_tol = self._float_tol return other @abstractmethod def copy(self): """ All expectations sublclasses must implement a shallow copy. """ pass def done(self): return (self._future.done() or not self._awaited) and self._success def __bool__(self): return self.done() def __or__(self, other): return ArsdkWhenAnyExpectation([self, other]) def __and__(self, other): return ArsdkWhenAllExpectations([self, other]) def __rshift__(self, other): return ArsdkWhenSequenceExpectations([self, other]) __nonzero__ = __bool__
class Auth(object): def __init__(self, stream): self.stream = stream self.auth_task = None self.bind_func = None self.response_fut = None self.responses = None stream.credentials = {} # feature_mechanisms need this stream.register_plugin('feature_mechanisms') stream.register_handler( Callback('Auth', StanzaPath('auth'), self._handle_auth)) stream.register_handler( Callback('Auth Response', StanzaPath('response'), self._handle_response)) stream.register_handler( Callback('Auth Abort', StanzaPath('abort'), self._handle_abort)) if LegacyAuth.available(self): # need to explicitly specify module here since xep_0078 # is purposefully unavailable by default stream.register_plugin('xep_0078', module=xep_0078) stream.register_handler( Callback('LegacyAuth', StanzaPath('iq/auth'), self._handle_legacy_auth)) stream.register_plugin('feature_bind') stream.register_handler( Callback('Bind', StanzaPath('iq/bind'), self._handle_bind)) stream.register_plugin('feature_session') stream.register_handler( Callback('Session', StanzaPath('iq/session'), self._handle_session)) stream.add_event_handler('disconnected', self._disconnected) async def get_features(self): features = StreamFeatures() available = await get_sasl_available(self) features['mechanisms'] = [m.name for m in available] if await LegacyAuth.available(self): features._get_plugin('auth') features._get_plugin('register') return features async def generate_resource_id(self): return uuid.uuid4().hex async def generate_anonymous_user(self): return uuid.uuid4().hex async def check_password(self, username, password): if not password: # client didn't supply a password, maybe they # want web session authentication web_user = self.stream.web_user if web_user: return await self.stream.auth_hook.check_webuser( self.stream, web_user, username) # if web authentication isn't available, # we won't allow passwordless logins elif password.startswith('//jid/'): # seems to be a session token return await self.stream.auth_hook.check_token( self.stream, username, password[6:]) elif settings.ALLOW_PLAIN_PASSWORD: # ordinary password return await self.stream.auth_hook.check_password( self.stream, username, password) return False async def prebound(self): if 'mechanisms' not in self.stream.features: await self._auth_success() await self._bind_attempt() await self._bind_success() def _disconnected(self, reason): self._abort_auth() async def _auth_success(self): jid = self.stream.boundjid.bare self.stream.update_logger({'jid': jid}) self.stream.logger.info('Authenticated as %s', jid) self.stream.features.add('mechanisms') self.stream.prepare_features() await self.stream.auth_hook.bind(self.stream) self.stream.event('auth_success', self.stream.boundjid) async def _async_challenge(self, data): challenge = auth_stanza.Challenge() challenge['value'] = data self.stream.send(challenge) while not self.responses: self.stream.send_thaw() await wrap_future(self.response_fut) self.response_fut = Future() self.stream.send_freeze() return self.responses.pop(0) async def _auth_task(self, auth, process): try: ret = process(auth) if iscoroutine(ret): await ret except Exception as e: self.stream.logger.info('Authentication failure') reply = auth_stanza.Failure() if isinstance(e, XMPPError): reply['condition'] = e.condition elif isinstance(e, UnicodeDecodeError): reply['condition'] = 'malformed-request' else: reply['condition'] = 'temporary-auth-failure' self.stream.logger.exception( 'Unhandled authentication exception') self.stream.send(reply) self.stream.send_thaw() else: # TODO: what if self.responses is non-empty # (i.e., we've received too many responses?) self.responses = None self.response_fut = None self.stream.send(auth_stanza.Success()) await self._auth_success() if self.bind_func: await self.bind_func self.stream.send_thaw() self.auth_task = None def _handle_auth(self, auth): # TODO: handle unexpected auth if 'mechanisms' in self.stream.features or self.auth_task: # already authenticated, or in progress # TODO: send error? return mech = get_sasl_by_name(auth['mechanism']) if not mech or not mech.available(self): reply = auth_stanza.Failure() reply['condition'] = 'invalid-mechanism' self.stream.send(reply) return self.stream.send_freeze() task = self._auth_task(auth, mech(self).process) self.auth_task = self.stream.loop.create_task(task) self.responses = [] self.response_fut = Future() def _handle_response(self, response): # TODO: handle unexpected response if not self.response_fut: # TODO: send error? return self.responses.append(response['value']) # wake up the thread, in case it's waiting for this try: self.response_fut.set_result() except InvalidStateError: # guess it wasn't waiting... it'll get to the # queued response on its own, eventually pass def _abort_auth(self): if self.auth_task: self.bind_func = None # auth_task and response_fut run in different threads, # so have to cancel both to avoid stuck threads self.auth_task.cancel() if self.response_fut: self.response_fut.cancel() self.responses = None self.response_fut = None self.stream.send_thaw() self.auth_task = None def _handle_abort(self, abort): self._abort_auth() reply = auth_stanza.Failure() reply['condition'] = 'aborted' self.stream.send(reply) async def _bind_attempt(self): # If there's a resource conflict, force a new resource ID. while not await self.stream.session_hook.bind(self.stream): resource = await self.generate_resource_id() self.stream.boundjid.resource = resource async def _bind_success(self): jid = self.stream.boundjid.full self.stream.update_logger({'jid': jid}) self.stream.logger.info('Bound to %s', jid) self.stream.session_bind_event.set() await self.stream.bind() self.stream.event('session_bind', self.stream.boundjid) self.stream.event('session_start') async def _bind_task(self, bind): self.bind_func = None resource = bind['bind']['resource'] if resource == '': resource = await self.generate_resource_id() self.stream.boundjid.resource = resource try: await self._bind_attempt() except Exception as e: bind.exception(e) return reply = bind.reply() reply['bind']['jid'] = self.stream.boundjid.full reply.send() await self._bind_success() def _handle_bind(self, bind): if self.stream.session_bind_event.is_set() or \ self.bind_func: # already bound # TODO: send error? return if 'mechanisms' not in self.stream.features and \ not self.auth_task: # not authenticated yet # TODO: send error? return # If authentication is not complete yet, then this must # be a pipelined request, and should be held until after # authentication. Otherwise, schedule it now. self.bind_func = self._bind_task(bind) if not self.auth_task: self.stream.loop.create_task(self.bind_func) def _handle_session(self, session): # RFC 3921 session creation is obsolete, # so return success but otherwise do nothing session.reply().send() async def _legacy_auth_get(self, iq): reply = iq.reply(clear=False) auth = reply['auth'] auth.clear() auth._set_sub_text('username', keep=True) if not settings.ALLOW_WEBUSER_LOGIN: auth._set_sub_text('password', keep=True) auth._set_sub_text('resource', keep=True) reply.send() async def _legacy_auth_task(self, iq, process): try: ret = process(iq) if iscoroutine(ret): await ret await self._bind_attempt() except Exception as e: self.stream.logger.info('Authentication failure') reply = iq.reply() reply['type'] = 'error' if isinstance(e, XMPPError): reply['error']['condition'] = e.condition else: reply['error']['condition'] = 'internal-server-error' self.stream.logger.exception( 'Unhandled authentication exception') reply.send() else: iq.reply().send() await self._auth_success() await self._bind_success() self.auth_task = None def _handle_legacy_auth(self, iq): # XEP-0078 if 'mechanisms' in self.stream.features or self.auth_task: # already authenticated, or in progress # TODO: send error? return type = iq['type'] auth = iq['auth'] if type == 'get': task = self._legacy_auth_get(auth) self.stream.loop.create_task(task) elif type == 'set': process = LegacyAuth(self).process task = self._legacy_auth_task(auth, process) self.auth_task = self.stream.loop.create_task(task) else: iq.unhandled()
def test_futures(self): f = Future() self.assertEqual(f.done(), False) self.assertEqual(f.running(), False) self.assertTrue(f.cancel()) self.assertTrue(f.cancelled()) with self.assertRaises(CancelledError): f.result() with self.assertRaises(CancelledError): f.exception() f = Future() f.set_running_or_notify_cancel() with self.assertRaises(TimeoutError): f.result(0.1) with self.assertRaises(TimeoutError): f.exception(0.1) f = Future() f.set_running_or_notify_cancel() f.set_result("result") self.assertEqual(f.result(), "result") self.assertEqual(f.exception(), None) f = Future() f.set_running_or_notify_cancel() f.set_exception(Exception("foo")) with self.assertRaises(Exception): f.result() class Ref(): def __init__(self, ref): self.ref = ref def set(self, ref): self.ref = ref # Test that done callbacks are called. called = Ref(False) f = Future() f.add_done_callback(lambda f: called.set(True)) f.set_result(None) self.assertTrue(called.ref) # Test that callbacks are called when cancelled. called = Ref(False) f = Future() f.add_done_callback(lambda f: called.set(True)) f.cancel() self.assertTrue(called.ref) # Test that callbacks are called immediately when the future is # already done. called = Ref(False) f = Future() f.set_result(None) f.add_done_callback(lambda f: called.set(True)) self.assertTrue(called.ref) count = Ref(0) f = Future() f.add_done_callback(lambda f: count.set(count.ref + 1)) f.add_done_callback(lambda f: count.set(count.ref + 1)) f.set_result(None) self.assertEqual(count.ref, 2) # Test that the callbacks are called with the future as argument. done_future = Ref(None) f = Future() f.add_done_callback(lambda f: done_future.set(f)) f.set_result(None) self.assertIs(f, done_future.ref)