async def wait_for_task(task_id: str) -> bool:
# Wait for a specific task
task = TaskState(task_id)
try:
await task.join()
return True
except TaskNotFoundException:
return False
async def test_task_from_service(rabbitmq_container, amqp_worker,
container_requester):
async with container_requester as requester:
resp, _ = await requester('GET', '/db/guillotina/@foobar')
state = TaskState(resp['task_id'])
await state.join(0.01)
assert await state.get_result() == 3
await asyncio.sleep(0.1) # prevent possible race condition here
assert amqp_worker.total_run == 1
async def handle_queued_job(self, channel, body, envelope, properties):
"""Callback triggered when there is a new job in the job channel (e.g:
a new task in a rabbitmq queue)
"""
logger.debug(f'Queued job {body}')
# Deserialize job description
if not isinstance(body, str):
body = body.decode('utf-8')
data = json.loads(body)
# Set status to scheduled
task_id = data['task_id']
dotted_name = data['func']
await update_task_scheduled(self.state_manager, task_id, eventlog=[])
logger.info(f'Received task: {task_id}: {dotted_name}')
self.measure_running_jobs(len(self._running))
# Block if we reached maximum number of running tasks
while len(self._running) >= self._max_running:
logger.info(f'Max running tasks reached: {self._max_running}')
await asyncio.sleep(self.sleep_interval)
self.last_activity = time.time()
# Create job object
self.last_activity = time.time()
job = Job(self.request, data, channel, envelope)
# Get the redis lock on the task so no other worker takes it
_id = job.data['task_id']
ts = TaskState(_id)
# Cancelation
if await ts.is_canceled():
logger.warning(f'Task {_id} has already been canceled')
raise TaskAlreadyCanceled(_id)
try:
await ts.acquire()
except TaskAlreadyAcquired:
logger.warning(f'Task {_id} is already running in another worker')
# TODO: ACK task
# Record job's data into global state
await self.state_manager.update(task_id, {
'job_data': job.data,
})
# Add the task to the loop and start it
task = self.loop.create_task(job())
task._job = job
job.task = task
self._running.append(task)
task.add_done_callback(self._task_done_callback)
async def handle_queued_job(self, channel, body, envelope, properties):
"""Callback triggered when there is a new job in the job channel (e.g:
a new task in a rabbitmq queue)
"""
logger.debug(f"Queued job {body}")
# Deserialize job description
if not isinstance(body, str):
body = body.decode("utf-8")
data = json.loads(body)
# Set status to scheduled
task_id = data["task_id"]
dotted_name = data["func"]
await update_task_scheduled(self.state_manager, task_id, eventlog=[])
logger.info(f"Received task: {task_id}: {dotted_name}")
self.measure_running_jobs(len(self._running))
# Block if we reached maximum number of running tasks
while len(self._running) >= self._max_running:
logger.info(f"Max running tasks reached: {self._max_running}")
await asyncio.sleep(self.sleep_interval)
self.last_activity = time.time()
# Create job object
self.last_activity = time.time()
job = Job(self.request, data, channel, envelope)
# Get the redis lock on the task so no other worker takes it
_id = job.data["task_id"]
ts = TaskState(_id)
# Cancelation
if await ts.is_canceled():
logger.warning(f"Task {_id} has already been canceled")
# Ack so that canceled job is removed from main queue
await channel.basic_client_ack(delivery_tag=envelope.delivery_tag)
return
if not await ts.acquire():
logger.warning(f"Task {_id} is already running in another worker")
await channel.basic_client_ack(delivery_tag=envelope.delivery_tag)
return
# Record job's data into global state
await self.state_manager.update(task_id, {"job_data": job.data})
# Add the task to the loop and start it
task = self.loop.create_task(job())
task._job = job
job.task = task
self._running.append(task)
task.add_done_callback(self._task_done_callback)
async def test_task_commits_data_from_service(rabbitmq_container, amqp_worker,
container_requester):
async with container_requester as requester:
await requester('POST',
'/db/guillotina',
data=json.dumps({
'@type': 'Item',
'id': 'foobar',
'title': 'blah'
}))
resp, _ = await requester('GET', '/db/guillotina/foobar/@foobar-write')
state = TaskState(resp['task_id'])
await state.join(0.01)
assert await state.get_result() == 'done!'
await asyncio.sleep(0.1) # prevent possible race condition here
assert amqp_worker.total_run == 1
resp, status = await requester('GET', '/db/guillotina/foobar')
assert resp['title'] == 'Foobar written'
async def test_task_commits_data_from_service(rabbitmq_container, amqp_worker,
container_requester):
async with container_requester as requester:
await requester(
"POST",
"/db/guillotina",
data=json.dumps({
"@type": "Item",
"id": "foobar",
"title": "blah"
}),
)
resp, _ = await requester("GET", "/db/guillotina/foobar/@foobar-write")
state = TaskState(resp["task_id"])
await state.join(0.01)
assert await state.get_result() == "done!"
await asyncio.sleep(0.1) # prevent possible race condition here
assert amqp_worker.total_run == 1
resp, status = await requester("GET", "/db/guillotina/foobar")
assert resp["title"] == "Foobar written"
async def test_async_gen_task_commits_data_from_service(
configured_state_manager, rabbitmq_container, amqp_worker,
container_requester):
async with container_requester as requester:
await requester('POST',
'/db/guillotina',
data=json.dumps({
'@type': 'Item',
'id': 'foobar',
'title': 'blah'
}))
resp, _ = await requester(
'GET', '/db/guillotina/foobar/@foobar-write-async-gen')
state = TaskState(resp['task_id'])
await state.join(0.01)
assert await state.get_status() == 'finished'
task_state = await state.get_state()
assert len(task_state['eventlog']) == 3
assert await state.get_result() == ['one', 'two', 'final']
await asyncio.sleep(0.1) # prevent possible race condition here
assert amqp_worker.total_run == 1
resp, status = await requester('GET', '/db/guillotina/foobar')
assert resp['title'] == 'CHANGED!'
async def update_status(self):
"""Updates status for running tasks and kills running tasks that have
been canceled.
"""
while True:
await asyncio.sleep(self.update_status_interval)
for task in self._running:
_id = task._job.data['task_id']
ts = TaskState(_id)
# Still working on the job: refresh task lock
await ts.refresh_lock()
# Cancel local tasks that have been cancelled in global
# state manager
async for val in self.state_manager.cancelation_list():
for task in self._running:
_id = task._job.data['task_id']
if _id == val:
logger.warning(f"Canceling task {_id}")
task.cancel()
self._running.remove(task)
await self._state_manager.clean_canceled(_id)
async def add_task(func,
*args,
_request=None,
_retries=3,
_task_id=None,
**kwargs):
"""Given a function and its arguments, it adds it as a task to be ran
by workers.
"""
# Get the request and prepare request data
if _request is None:
_request = get_current_request()
req_data: SerializedRequest = serialize_request(_request)
if _task_id is None:
task_id = generate_task_id()
else:
task_id = _task_id
dotted_name = get_dotted_name(func)
retries = 0
while True:
# Get the rabbitmq connection
try:
channel, transport, protocol = await amqp.get_connection()
except AMQPConfigurationNotFoundError:
logger.warning(
f"Could not schedule {dotted_name}, AMQP settings not configured"
)
return
try:
state = TaskState(task_id)
db = task_vars.db.get()
container = task_vars.container.get()
logger.info(f"Scheduling task: {task_id}: {dotted_name}")
data = json.dumps({
"func": dotted_name,
"args": args,
"kwargs": kwargs,
"db_id": getattr(db, "id", None),
"container_id": getattr(container, "id", None),
"req_data": req_data,
"task_id": task_id,
})
# Publish task data on rabbitmq
await channel.publish(
data,
exchange_name=app_settings["amqp"]["exchange"],
routing_key=app_settings["amqp"]["queue"],
properties={"delivery_mode": 2},
)
# Update tasks's global state
state_manager = get_state_manager()
await update_task_scheduled(state_manager,
task_id,
updated=time.time())
logger.info(f"Scheduled task: {task_id}: {dotted_name}")
return state
except (aioamqp.AmqpClosedConnection,
aioamqp.exceptions.ChannelClosed):
await amqp.remove_connection()
if retries >= _retries:
raise
retries += 1
async def add_task(func,
*args,
_request=None,
_retries=3,
_task_id=None,
**kwargs):
"""Given a function and its arguments, it adds it as a task to be ran
by workers.
"""
# Get the request and prepare request data
if _request is None:
_request = get_current_request()
req_data = {
'url': str(_request.url),
'headers': dict(_request.headers),
'method': _request.method,
'annotations': getattr(_request, 'annotations', {})
}
user = get_authenticated_user()
if user is not None:
try:
req_data['user'] = {
'id':
user.id,
'roles': [
name for name, setting in user.roles.items()
if setting == Allow
],
'groups':
user.groups,
'headers':
dict(_request.headers),
'data':
getattr(user, 'data', {})
}
except AttributeError:
pass
container = task_vars.container.get()
if container is not None:
req_data['container_url'] = IAbsoluteURL(container, _request)()
if _task_id is None:
task_id = generate_task_id()
else:
task_id = _task_id
retries = 0
while True:
# Get the rabbitmq connection
channel, transport, protocol = await amqp.get_connection()
try:
state = TaskState(task_id)
dotted_name = get_dotted_name(func)
db = task_vars.db.get()
logger.info(f'Scheduling task: {task_id}: {dotted_name}')
data = json.dumps({
'func': dotted_name,
'args': args,
'kwargs': kwargs,
'db_id': getattr(db, 'id', None),
'container_id': getattr(container, 'id', None),
'req_data': req_data,
'task_id': task_id
})
# Publish task data on rabbitmq
await channel.publish(
data,
exchange_name=app_settings['amqp']['exchange'],
routing_key=app_settings['amqp']['queue'],
properties={'delivery_mode': 2})
# Update tasks's global state
state_manager = get_state_manager()
await update_task_scheduled(state_manager,
task_id,
updated=time.time())
logger.info(f'Scheduled task: {task_id}: {dotted_name}')
return state
except (aioamqp.AmqpClosedConnection,
aioamqp.exceptions.ChannelClosed):
await amqp.remove_connection()
if retries >= _retries:
raise
retries += 1