27 lines
918 B
Python
27 lines
918 B
Python
import asyncio
|
|
from asyncio import Task
|
|
from typing import List
|
|
|
|
from fastapi import FastAPI
|
|
from loguru import logger
|
|
|
|
|
|
async def create_task(app: FastAPI, identifier, callable, *args, **kwargs):
|
|
async with app.state.tasks_lock:
|
|
task = asyncio.create_task(callable(*args, **kwargs))
|
|
app.state.tasks[identifier].append(task)
|
|
logger.info(f"Task {identifier} created")
|
|
|
|
|
|
async def cancel_task(app: FastAPI, identifier):
|
|
async with app.state.tasks_lock:
|
|
tasks: List[Task] = app.state.tasks.get(identifier)
|
|
if not tasks:
|
|
logger.info(f"Task {identifier} not found, can't cancel")
|
|
for task in tasks:
|
|
try:
|
|
task.cancel()
|
|
await task # Wait for the task to finish cancellation
|
|
except asyncio.CancelledError:
|
|
logger.info(f"Task {identifier} cancelled")
|
|
del app.state.tasks[identifier] |