feat: Implement cancellation of stale generations in the service and repository, along with a new test.

This commit is contained in:
xds
2026-02-13 17:30:11 +03:00
parent 30138bab38
commit 279cb5c6f6
6 changed files with 109 additions and 6 deletions

32
aiws.py
View File

@@ -120,6 +120,17 @@ assets_router.message.middleware(AuthMiddleware(repo=users_repo, admin_id=ADMIN_
gen_router.message.middleware(AlbumMiddleware(latency=0.8))
async def start_scheduler(service: GenerationService):
while True:
try:
logger.info("Running scheduler for stacked generation killing")
await service.cleanup_stale_generations()
except asyncio.CancelledError:
break
except Exception as e:
logger.error(f"Scheduler error: {e}")
await asyncio.sleep(300) # Check every 5 minutes
# --- LIFESPAN (Запуск FastAPI + Bot) ---
@asynccontextmanager
async def lifespan(app: FastAPI):
@@ -151,17 +162,28 @@ async def lifespan(app: FastAPI):
# )
# print("🤖 Bot polling started")
# 3. ЗАПУСК ШЕДУЛЕРА
scheduler_task = asyncio.create_task(start_scheduler(generation_service))
print("⏰ Scheduler started")
yield
# --- SHUTDOWN ---
print("🛑 Shutting down...")
# 3. Остановка бота
polling_task.cancel()
# 4. Остановка шедулера
scheduler_task.cancel()
try:
await polling_task
await scheduler_task
except asyncio.CancelledError:
print("🤖 Bot polling stopped")
print("⏰ Scheduler stopped")
# 3. Остановка бота
# polling_task.cancel()
# try:
# await polling_task
# except asyncio.CancelledError:
# print("🤖 Bot polling stopped")
# 4. Отключение БД
# Обычно Motor закрывать не обязательно при выходе, но хорошим тоном считается