161 lines
6.5 KiB
Python
161 lines
6.5 KiB
Python
"""
|
||
Декораторы для асинхронной обработки запросов пользователей
|
||
"""
|
||
import asyncio
|
||
import functools
|
||
from typing import Callable, Any
|
||
from aiogram import types
|
||
from task_manager import task_manager, TaskPriority
|
||
import uuid
|
||
import logging
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def async_user_action(priority: TaskPriority = TaskPriority.NORMAL, timeout: float = 30.0):
|
||
"""
|
||
Декоратор для асинхронной обработки действий пользователей
|
||
|
||
Args:
|
||
priority: Приоритет задачи
|
||
timeout: Таймаут выполнения в секундах
|
||
"""
|
||
def decorator(func: Callable) -> Callable:
|
||
@functools.wraps(func)
|
||
async def wrapper(*args, **kwargs):
|
||
# Извлекаем информацию о пользователе
|
||
user_id = None
|
||
action_name = func.__name__
|
||
|
||
# Ищем пользователя в аргументах
|
||
for arg in args:
|
||
if isinstance(arg, (types.Message, types.CallbackQuery)):
|
||
user_id = arg.from_user.id
|
||
break
|
||
|
||
if user_id is None:
|
||
# Если не нашли пользователя, выполняем синхронно
|
||
logger.warning(f"Не удалось определить user_id для {action_name}, выполнение синхронно")
|
||
return await func(*args, **kwargs)
|
||
|
||
# Генерируем ID задачи
|
||
task_id = f"{action_name}_{user_id}_{uuid.uuid4().hex[:8]}"
|
||
|
||
try:
|
||
# Добавляем задачу в очередь
|
||
await task_manager.add_task(
|
||
task_id,
|
||
user_id,
|
||
func,
|
||
*args,
|
||
priority=priority,
|
||
timeout=timeout,
|
||
**kwargs
|
||
)
|
||
|
||
logger.debug(f"Задача {task_id} добавлена в очередь для пользователя {user_id}")
|
||
|
||
except ValueError as e:
|
||
# Превышен лимит задач пользователя
|
||
logger.warning(f"Лимит задач для пользователя {user_id}: {e}")
|
||
|
||
# Отправляем сообщение о превышении лимита
|
||
if isinstance(args[0], types.Message):
|
||
message = args[0]
|
||
await message.answer(
|
||
"⚠️ Вы превысили лимит одновременных запросов. "
|
||
"Пожалуйста, дождитесь завершения предыдущих операций."
|
||
)
|
||
elif isinstance(args[0], types.CallbackQuery):
|
||
callback = args[0]
|
||
await callback.answer(
|
||
"⚠️ Превышен лимит запросов. Дождитесь завершения предыдущих операций.",
|
||
show_alert=True
|
||
)
|
||
|
||
return None
|
||
|
||
return wrapper
|
||
return decorator
|
||
|
||
|
||
def admin_async_action(priority: TaskPriority = TaskPriority.HIGH, timeout: float = 60.0):
|
||
"""
|
||
Декоратор для асинхронной обработки действий администраторов
|
||
(повышенный приоритет и больший таймаут)
|
||
"""
|
||
return async_user_action(priority=priority, timeout=timeout)
|
||
|
||
|
||
def critical_action(timeout: float = 120.0):
|
||
"""
|
||
Декоратор для критических действий (розыгрыши, важные операции)
|
||
"""
|
||
return async_user_action(priority=TaskPriority.CRITICAL, timeout=timeout)
|
||
|
||
|
||
def db_operation(timeout: float = 15.0):
|
||
"""
|
||
Декоратор для операций с базой данных
|
||
"""
|
||
return async_user_action(priority=TaskPriority.NORMAL, timeout=timeout)
|
||
|
||
|
||
# Функции для работы со статистикой задач
|
||
|
||
async def get_task_stats() -> dict:
|
||
"""Получить общую статистику задач"""
|
||
return task_manager.get_stats()
|
||
|
||
|
||
async def get_user_task_info(user_id: int) -> dict:
|
||
"""Получить информацию о задачах пользователя"""
|
||
return task_manager.get_user_stats(user_id)
|
||
|
||
|
||
async def format_task_stats() -> str:
|
||
"""Форматированная статистика для админов"""
|
||
stats = await get_task_stats()
|
||
|
||
text = "📊 **Статистика обработки задач:**\n\n"
|
||
text += f"🟢 Активных воркеров: {stats['workers_count']}\n"
|
||
text += f"⚙️ Выполняется задач: {stats['active_tasks']}\n"
|
||
text += f"📋 В очереди: {stats['queue_size']}\n"
|
||
text += f"✅ Выполнено: {stats['completed_tasks']}\n"
|
||
text += f"❌ Ошибок: {stats['failed_tasks']}\n\n"
|
||
|
||
if stats['user_tasks']:
|
||
text += "👥 **Активные пользователи:**\n"
|
||
for user_id, task_count in stats['user_tasks'].items():
|
||
if task_count > 0:
|
||
text += f"• ID {user_id}: {task_count} задач\n"
|
||
|
||
return text
|
||
|
||
|
||
# Middleware для автоматического управления задачами
|
||
|
||
class TaskManagerMiddleware:
|
||
"""Middleware для управления менеджером задач"""
|
||
|
||
def __init__(self):
|
||
self.started = False
|
||
|
||
async def __call__(self, handler: Callable, event: types.TelegramObject, data: dict):
|
||
# Запускаем менеджер при первом обращении
|
||
if not self.started:
|
||
await task_manager.start()
|
||
self.started = True
|
||
logger.info("Менеджер задач запущен через middleware")
|
||
|
||
# Продолжаем обработку
|
||
return await handler(event, data)
|
||
|
||
|
||
# Функция для изящного завершения
|
||
|
||
async def shutdown_task_manager():
|
||
"""Завершение работы менеджера задач"""
|
||
logger.info("Завершение работы менеджера задач...")
|
||
await task_manager.stop()
|
||
logger.info("Менеджер задач остановлен") |