def test_context_manager_cancel(self):
tq = tasks.CompletionQueue()
event = locks.Event()
t1 = self.k.spawn(event.wait)
tq.put_nonblocking(t1)
t2 = self.k.spawn(event.wait)
tq.put_nonblocking(t2)
t3 = self.k.spawn(raises('test message'))
tq.put_nonblocking(t3)
async def do_with_queue():
async with tq:
raise Exception('some error')
with self.assertRaisesRegex(Exception, r'some error'):
self.k.run(do_with_queue)
self.assertTrue(tq.is_closed())
self.assertFalse(tq)
self.assertEqual(len(tq), 0)
for t in (t1, t2):
self.assertTrue(t.is_completed())
with self.assertRaises(errors.Cancelled):
t.get_result_nonblocking()
self.assertTrue(t3.is_completed())
with self.assertRaisesRegex(Exception, r'test message'):
t3.get_result_nonblocking()
def test_get_nonblocking(self):
tq = tasks.CompletionQueue()
with self.assertRaises(tasks.Empty):
tq.get_nonblocking()
gettable_task = self.k.spawn(tq.gettable())
with self.assertRaises(errors.KernelTimeout):
self.k.run(timeout=0.01)
self.assertFalse(gettable_task.is_completed())
event = locks.Event()
t1 = self.k.spawn(event.wait)
t2 = self.k.spawn(event.wait)
tq.put_nonblocking(t1)
tq.put_nonblocking(t2)
tq.close()
with self.assertRaises(errors.KernelTimeout):
self.k.run(timeout=0.01)
with self.assertRaises(tasks.Empty):
tq.get_nonblocking()
event.set()
self.k.run(timeout=0.01)
self.assertTrue(gettable_task.is_completed())
self.assertCountEqual(
[tq.get_nonblocking(), tq.get_nonblocking()],
[t1, t2],
)
with self.assertRaises(tasks.Closed):
tq.get_nonblocking()
self.k.run(tq.gettable())
def test_spawn(self):
tq = tasks.CompletionQueue()
tq.close()
self.assertEqual(self.k.get_all_tasks(), [])
with self.assertRaises(tasks.Closed):
tq.spawn(square)
self.assertEqual(self.k.get_all_tasks(), [])
def setUp(self):
super().setUp()
self.main_task = None
self.agent_queue = tasks.CompletionQueue()
self.graceful_exit = locks.Event()
self.signal_queue = queues.Queue()
mock = unittest.mock.patch(agents.__name__ + '.signals').start()
mock.SignalSource().__enter__().get = self.signal_queue.get
self._assert_logs = self.assertLogs(agents.__name__, level='DEBUG')
self.cm = self._assert_logs.__enter__()
async def run_servers(url, num_ctxs):
with contextlib.ExitStack() as stack:
socket = stack.enter_context(nng.asyncs.Socket(nng.Protocols.REP0))
socket.listen(url)
async with tasks.CompletionQueue() as servers:
for _ in range(num_ctxs):
ctx = stack.enter_context(nng.asyncs.Context(socket))
servers.spawn(serve(ctx))
servers.close()
async for server in servers:
server.get_result_nonblocking()
async def serve(self):
LOG.debug('start server: %r', self._socket)
with self._socket:
if self._max_connections <= 0:
capacity = self._max_connections
else:
# +1 for the `_accept` task.
capacity = self._max_connections + 1
async with tasks.CompletionQueue(capacity) as queue:
await servers.supervise_server(
queue,
(queue.spawn(self._accept(queue)), ),
)
LOG.debug('stop server: %r', self._socket)
def __init__(self, engine, publisher):
self._engine = engine
self._manager = connections.ConnectionManager(self._engine.connect())
self._metadata = sqlalchemy.MetaData()
self._tables = schemas.make_tables(self._metadata)
self._tx_revision = None
# A transaction is automatically rolled back if it is inactive
# after a certain amount of time. This is a fail-safe mechanism
# to prevent deadlocks due to client crashes.
self._timer_queue = tasks.CompletionQueue()
self._tx_expiration = time.monotonic()
# For publishing database events.
self._publisher = publisher
self._pending_events = collections.deque()
def test_queue(self):
tq = tasks.CompletionQueue()
self.assertFalse(tq.is_full())
self.assertFalse(tq.is_closed())
self.assertFalse(tq)
self.assertEqual(len(tq), 0)
t1 = self.k.spawn(square(1))
tq.put_nonblocking(t1)
self.assertFalse(tq.is_full())
self.assertFalse(tq.is_closed())
self.assertTrue(tq)
self.assertEqual(len(tq), 1)
t2 = self.k.spawn(square(1))
tq.put_nonblocking(t2)
self.assertFalse(tq.is_full())
self.assertFalse(tq.is_closed())
self.assertTrue(tq)
self.assertEqual(len(tq), 2)
tq.close()
self.assertFalse(tq.is_full())
self.assertTrue(tq.is_closed())
self.assertTrue(tq)
self.assertEqual(len(tq), 2)
with self.assertRaises(tasks.Closed):
tq.put_nonblocking(None)
ts = set()
ts.add(self.k.run(tq.get, timeout=1))
self.assertFalse(tq.is_full())
self.assertTrue(tq.is_closed())
self.assertTrue(tq)
self.assertEqual(len(tq), 1)
ts.add(self.k.run(tq.get, timeout=1))
self.assertFalse(tq.is_full())
self.assertTrue(tq.is_closed())
self.assertFalse(tq)
self.assertEqual(len(tq), 0)
with self.assertRaises(tasks.Closed):
self.k.run(tq.get)
self.assertEqual(ts, {t1, t2})
def test_completion_queue(self):
f1 = futures.Future()
f2 = futures.Future()
f3 = futures.Future()
queue = tasks.CompletionQueue()
queue.put_nonblocking(f1)
queue.put_nonblocking(f2)
queue.put_nonblocking(f3)
with self.assertRaises(tasks.Empty):
queue.get_nonblocking()
f2.set_result(None)
self.assertIs(queue.get_nonblocking(), f2)
f3.set_result(None)
self.assertIs(queue.get_nonblocking(), f3)
f1.set_result(None)
self.assertIs(queue.get_nonblocking(), f1)
def test_close_repeatedly(self):
tq = tasks.CompletionQueue()
self.assertFalse(tq.is_closed())
t1 = self.k.spawn(square(1))
tq.put_nonblocking(t1)
self.assertEqual(tq.close(True), [])
self.assertTrue(tq.is_closed())
self.assertEqual(tq.close(False), [t1])
self.assertTrue(tq.is_closed())
self.assertEqual(tq.close(True), [])
self.assertTrue(tq.is_closed())
self.assertEqual(tq.close(False), [])
self.assertTrue(tq.is_closed())
self.assertFalse(t1.is_completed())
self.k.run(timeout=1)
self.assertTrue(t1.is_completed())
def test_async_iterator(self):
tq = tasks.CompletionQueue()
expect = {
tq.spawn(square(1)),
tq.spawn(square(2)),
tq.spawn(square(3)),
}
tq.close()
async def async_iter():
actual = set()
async for task in tq:
actual.add(task)
return actual
self.assertEqual(
self.k.run(async_iter, timeout=1),
expect,
)
def test_queue(self):
with self.assertRaises(LookupError):
contexts.get_kernel()
tq = tasks.CompletionQueue()
self.assertFalse(tq.is_closed())
self.assertFalse(tq)
self.assertEqual(len(tq), 0)
task = Task(None, square(7))
tq.put_nonblocking(task)
self.assertFalse(tq.is_closed())
self.assertTrue(tq)
self.assertEqual(len(tq), 1)
tq.close()
self.assertTrue(tq.is_closed())
self.assertTrue(tq)
self.assertEqual(len(tq), 1)
self.assertIsNone(task.tick(None, None))
def test_context_manager(self):
tq = tasks.CompletionQueue()
t1 = self.k.spawn(square(1))
tq.put_nonblocking(t1)
t2 = self.k.spawn(square(2))
tq.put_nonblocking(t2)
async def do_with_queue():
async with tq:
return 42
self.assertEqual(self.k.run(do_with_queue, timeout=1), 42)
self.assertTrue(tq.is_closed())
self.assertFalse(tq)
self.assertEqual(len(tq), 0)
for t, x in [(t1, 1), (t2, 2)]:
self.assertTrue(t.is_completed())
self.assertEqual(t.get_result_nonblocking(), x * x)
def test_put_and_capacity(self):
tq = tasks.CompletionQueue(capacity=1)
self.assertFalse(tq.is_full())
event = locks.Event()
t1 = self.k.spawn(event.wait)
t2 = self.k.spawn(square(2))
tq.put_nonblocking(t1)
self.assertTrue(tq.is_full())
with self.assertRaises(tasks.Full):
tq.put_nonblocking(t2)
puttable_task = self.k.spawn(tq.puttable())
with self.assertRaises(errors.KernelTimeout):
self.k.run(timeout=0.01)
self.assertFalse(puttable_task.is_completed())
self.assertTrue(tq.is_full())
event.set()
self.k.run(timeout=0.01)
self.assertFalse(tq.is_full())
self.assertTrue(puttable_task.is_completed())
tq.close()
self.k.run(tq.puttable())
'supervise_agents',
'agent_queue',
'graceful_exit',
'grace_period',
# shutdown_agents.
'shutdown_queue',
)
PARAMS = parameters.define(
agents.__name__,
parameters.Namespace(grace_period=parameters.Parameter(4,
type=(int, float),
unit='seconds'), ),
)
startup.set(LABELS.agent_queue, tasks.CompletionQueue())
startup.set(LABELS.graceful_exit, locks.Event())
startup.set(LABELS.shutdown_queue, queues.Queue())
utils.depend_parameter_for(LABELS.grace_period, PARAMS.grace_period)
utils.define_binder(
agents.supervise_agents,
LABELS.supervise_agents,
{
'agent_queue': LABELS.agent_queue,
'graceful_exit': LABELS.graceful_exit,
'grace_period': LABELS.grace_period,
},
)
def __init__(self, handler):
self._handler = handler
self._handler_queue = tasks.CompletionQueue()
def setUp(self):
super().setUp()
self.queue = tasks.CompletionQueue()
self._assert_logs = self.assertLogs(servers.__name__, level='DEBUG')
self.cm = self._assert_logs.__enter__()
def __init__(self):
self.queue = tasks.CompletionQueue()
self._leftover = set()