def test_disallow_repeated_enable(self):
with signals.SignalSource() as source:
source.enable(0)
with self.assertRaises(AssertionError):
source.enable(0)
source.disable(0)
with self.assertRaises(AssertionError):
source.disable(0)
async def handle_signals(duration):
print('pid: %d' % os.getpid())
timers.timeout_after(duration)
with signals.SignalSource() as source:
signums = [signal.SIGINT, signal.SIGTERM]
print('handle signals for %.3f seconds: %r' % (duration, signums))
for signum in signums:
source.enable(signum)
try:
while True:
print('receive signal: %r' % await source.get())
except timers.Timeout:
print('timeout')
def test_disallow_nested_use(self):
source = signals.SignalSource()
self.assertIsNone(source._wakeup_fd)
self.assertEqual(source._handlers, {})
with source:
with self.assertRaises(AssertionError):
with source:
pass
self.assertIsNone(source._wakeup_fd)
self.assertEqual(source._handlers, {})
# But consecutive use is fine.
with source:
pass
self.assertIsNone(source._wakeup_fd)
self.assertEqual(source._handlers, {})
self.signal_mock.signal.assert_not_called()
self.signal_mock.siginterrupt.assert_not_called()
def test_get(self):
with signals.SignalSource() as source:
source._sock_w.send(b'\x02')
kernels.run(source.get())
self.signal_mock.Signals.assert_called_once_with(2)
def test_singleton(self):
self.assertIs(signals.SignalSource(), signals.SignalSource())
async def supervise_agents(
agent_queue,
graceful_exit,
grace_period, # Unit: seconds.
):
"""Supervise agents.
The supervisor starts exiting when any of the following happens:
* An agent exits normally.
* An agent errs out (including being cancelled).
* A signal is delivered.
* The graceful_exit event is set.
During the exit:
* It closes the agent_queue; thus no new agent can be spawned.
* If it starts exiting due to an agent erring out, it cancels
remaining agents and exits.
* Else, it does a graceful exit:
* It sets the graceful_exit event.
* It waits for the grace_period. If during this period, an agent
errs out or another signal is delivered, it cancels remaining
agents and exits.
* After the grace_period, it cancels remaining agents and exits.
"""
main_task = tasks.get_current_task()
async with contextlib.AsyncExitStack() as stack:
def start_exiting():
if agent_queue.is_closed():
return False
agent_queue.close()
graceful_exit.set()
stack.enter_context(
timers.timeout_after(grace_period, task=main_task)
)
return True
signal_source = stack.enter_context(signals.SignalSource())
for signum in EXIT_SIGNUMS:
signal_source.enable(signum)
# Make sure that remaining agents are cancelled on error.
await stack.enter_async_context(agent_queue)
internal_tasks = [
tasks.spawn_onto_stack(
awaitable, stack, always_cancel=True, log_error=False
) for awaitable in (
join_agents(agent_queue, start_exiting),
request_graceful_exit(graceful_exit, start_exiting),
receive_signals(signal_source, start_exiting),
)
]
join_agents_task = internal_tasks[0]
try:
async for task in tasks.as_completed(internal_tasks):
task.get_result_nonblocking()
if task is join_agents_task:
# If all agents are completed, let the supervisor
# exit from here since we don't quite care whether
# another signal is delivered.
break
except timers.Timeout:
raise SupervisorError('grace period exceeded') from None