Skip to content

Commit

Permalink
WIP: Add graceful cancel to cancel scopes
Browse files Browse the repository at this point in the history
  • Loading branch information
nmalaguti committed Feb 20, 2019
1 parent 2025947 commit 7ca63d5
Showing 1 changed file with 89 additions and 6 deletions.
95 changes: 89 additions & 6 deletions trio/_core/_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -168,11 +168,13 @@ class CancelScope:
_has_been_entered = attr.ib(default=False, init=False)
_effective_deadline = attr.ib(default=inf, init=False)
cancel_called = attr.ib(default=False, init=False)
graceful_cancel_called = attr.ib(default=False, init=False)
cancelled_caught = attr.ib(default=False, init=False)

# Constructor arguments:
_deadline = attr.ib(default=inf, kw_only=True)
_shield = attr.ib(default=False, kw_only=True)
_graceful = attr.ib(default=None, kw_only=True)

@enable_ki_protection
def __enter__(self):
Expand Down Expand Up @@ -325,6 +327,19 @@ def shield(self, new_value):
for task in self._tasks:
task._attempt_delivery_of_any_pending_cancel()

@property
def graceful(self):
"""If this is set to :data:`True`, then the code inside this
scope will receive :exc:`~trio.Cancelled` exceptions from scopes
that are outside this scope that have been gracefully canceled.
If this is set to :data:`None` then this scope will inherit its
graceful cancel behavior from its parent scope.
Defaults to :data:`None`, though this can be overridden by the
``graceful=`` argument to the :class:`~trio.CancelScope` constructor.
"""
return self._graceful

@enable_ki_protection
def cancel(self):
"""Cancels this scope immediately.
Expand All @@ -339,6 +354,27 @@ def cancel(self):
for task in self._tasks:
task._attempt_delivery_of_any_pending_cancel()

@enable_ki_protection
def graceful_cancel(self, seconds=None):
"""Gracefully cancel this scope immediately.
Args:
seconds (float): Optional. If provided, this will change the deadline
of the current scope to be the number of seconds in the future.
Raises:
ValueError: if *seconds* is negative.
"""
if self.cancel_called or self.graceful_cancel_called:
return
if seconds is not None:
if seconds < 0:
raise ValueError("new deadline must be in the future")
self.deadline = current_time() + float(seconds)
self.graceful_cancel_called = True
for task in self._tasks:
task._attempt_delivery_of_any_pending_cancel()

def _add_task(self, task):
self._tasks.add(task)
task._cancel_stack.append(self)
Expand All @@ -365,10 +401,7 @@ def _exc_filter(self, exc):

def _close(self, exc):
scope_task = current_task()
if (
exc is not None and self.cancel_called
and scope_task._pending_cancel_scope() is self
):
if exc is not None and scope_task._pending_cancel_scope() is self:
exc = MultiError.filter(self._exc_filter, exc)
with self._might_change_effective_deadline():
self._remove_task(scope_task)
Expand Down Expand Up @@ -643,16 +676,51 @@ def __del__(self):
################################################################


def _pending_graceful_cancel_scope(cancel_stack):
# Return the outermost scope that is graceful, graceful_cancelled, and not outside a shield
pending_scope = None
graceful_cancel_called = False
scope_is_graceful = False
level = len(cancel_stack)

for i, scope in enumerate(cancel_stack):
if scope.shield:
pending_scope = None
level = len(cancel_stack)
graceful_cancel_called = False
scope_is_graceful = False
if scope.graceful is False: # None implies inherit
pending_scope = None
level = len(cancel_stack)
scope_is_graceful = False
if scope.graceful:
scope_is_graceful = True
if scope.graceful_cancel_called:
graceful_cancel_called = True
if pending_scope is None and scope_is_graceful and graceful_cancel_called:
level = i
pending_scope = scope

return level, pending_scope


def _pending_cancel_scope(cancel_stack):
# Return the outermost exception that is is not outside a shield.
pending_scope = None
for scope in cancel_stack:
level = len(cancel_stack)
for i, scope in enumerate(cancel_stack):
# Check shield before _exc, because shield should not block
# processing of *this* scope's exception
if scope.shield:
level = len(cancel_stack)
pending_scope = None
if pending_scope is None and scope.cancel_called:
level = i
pending_scope = scope

# Pick the outermost scope between cancel and graceful cancel
_, pending_scope = min((level, pending_scope), _pending_graceful_cancel_scope(cancel_stack))

return pending_scope


Expand Down Expand Up @@ -1671,12 +1739,27 @@ def current_effective_deadline():
"""
task = current_task()
deadline = inf
graceful_deadline = inf
graceful_cancel_called = False
scope_is_graceful = False
for scope in task._cancel_stack:
if scope._shield:
deadline = inf
graceful_deadline = inf
graceful_cancel_called = False
scope_is_graceful = False
if scope.graceful is False: # None implies inherit
graceful_deadline = inf
scope_is_graceful = False
if scope.graceful:
scope_is_graceful = True
if scope.graceful_cancel_called:
graceful_cancel_called = True
if scope.cancel_called:
deadline = -inf
deadline = min(deadline, scope._deadline)
if scope_is_graceful and graceful_cancel_called:
graceful_deadline = -inf
deadline = min(deadline, graceful_deadline, scope._deadline)
return deadline


Expand Down

0 comments on commit 7ca63d5

Please sign in to comment.