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}") if not task.done(): task.cancel() self._running.remove(task) await self._state_manager.clean_canceled(_id)
async def cancel_task(task_id): """It cancels a task by id. Returns wether it could be cancelled. """ task = TaskState(task_id) success = await task.cancel() return success
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 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