mirror of
https://github.com/daimyomizukagebay61/SteamPanel.git
synced 2026-07-25 11:44:28 +00:00
v 3.0.3
This commit is contained in:
Manchik
@@ -0,0 +1,249 @@
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any, Protocol
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from app.config import settings
|
||||
|
||||
DEFAULT_MAX_PARALLEL = 10
|
||||
DEFAULT_ACCOUNT_TIMEOUT = 90 # fallback for actions not in the map below
|
||||
|
||||
# Per-action timeouts (seconds). Fast read-only actions get 60s;
|
||||
# write actions that involve SMS/email prompts get 120s.
|
||||
TASK_TIMEOUTS: dict[str, int] = {
|
||||
"validate": 60,
|
||||
"remove_guard": 60,
|
||||
"change_password": 60,
|
||||
"random_password": 60,
|
||||
"change_phone": 120,
|
||||
"change_email": 120,
|
||||
}
|
||||
|
||||
|
||||
def _read_max_threads() -> int:
|
||||
from app.config import read_validation_settings
|
||||
return int(read_validation_settings().get("max_threads", DEFAULT_MAX_PARALLEL))
|
||||
|
||||
|
||||
def _read_account_timeout() -> int:
|
||||
from app.config import read_validation_settings
|
||||
return int(read_validation_settings().get("account_timeout", DEFAULT_ACCOUNT_TIMEOUT))
|
||||
|
||||
|
||||
class TaskHandler(Protocol):
|
||||
async def __call__(self, account: dict, params: dict, *, task_id: str) -> None: ...
|
||||
|
||||
|
||||
class TaskManager:
|
||||
def __init__(self) -> None:
|
||||
self._running: dict[str, asyncio.Task] = {}
|
||||
self._handlers: dict[str, TaskHandler] = {}
|
||||
self._prompts: dict[str, dict[str, Any]] = {}
|
||||
self._steps: dict[str, dict[str, Any]] = {}
|
||||
self._active_counts: dict[str, int] = {}
|
||||
self._account_results: dict[str, dict[int, dict[str, str]]] = {}
|
||||
self._account_steps: dict[str, dict[int, dict[str, int]]] = {}
|
||||
|
||||
def register_handler(self, action: str, handler: TaskHandler) -> None:
|
||||
self._handlers[action] = handler
|
||||
|
||||
async def submit(
|
||||
self,
|
||||
task_type: str,
|
||||
accounts: list[dict],
|
||||
params: dict,
|
||||
) -> str:
|
||||
task_id = uuid.uuid4().hex[:12]
|
||||
account_ids = [acc.get("id", 0) for acc in accounts]
|
||||
|
||||
from app.database import get_db
|
||||
|
||||
db = await get_db()
|
||||
await db.execute(
|
||||
"INSERT INTO tasks (id, type, status, total, account_ids) VALUES (?, ?, 'running', ?, ?)",
|
||||
(task_id, task_type, len(accounts), json.dumps(account_ids)),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async_task = asyncio.create_task(
|
||||
self._execute(task_id, task_type, accounts, params)
|
||||
)
|
||||
self._running[task_id] = async_task
|
||||
return task_id
|
||||
|
||||
async def _execute(
|
||||
self,
|
||||
task_id: str,
|
||||
task_type: str,
|
||||
accounts: list[dict],
|
||||
params: dict,
|
||||
) -> None:
|
||||
from app.database import get_db
|
||||
|
||||
handler = self._handlers.get(task_type)
|
||||
if not handler:
|
||||
db = await get_db()
|
||||
await db.execute(
|
||||
"UPDATE tasks SET status = 'failed', error = ? WHERE id = ?",
|
||||
(f"No handler for action: {task_type}", task_id),
|
||||
)
|
||||
await db.commit()
|
||||
return
|
||||
|
||||
completed = 0
|
||||
errors: list[str] = []
|
||||
lock = asyncio.Lock()
|
||||
max_parallel = _read_max_threads()
|
||||
account_timeout = TASK_TIMEOUTS.get(task_type, _read_account_timeout())
|
||||
sem = asyncio.Semaphore(max_parallel)
|
||||
self._active_counts[task_id] = 0
|
||||
self._account_results[task_id] = {}
|
||||
total = len(accounts)
|
||||
|
||||
async def process_one(account: dict, idx: int) -> None:
|
||||
nonlocal completed
|
||||
login = account.get("login", "?")
|
||||
acc_id = account.get("id", 0)
|
||||
logger.debug(f"[{task_type}] ({idx}/{total}) processing {login}...")
|
||||
|
||||
async with sem:
|
||||
async with lock:
|
||||
self._active_counts[task_id] = self._active_counts.get(task_id, 0) + 1
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
handler(account, params, task_id=task_id),
|
||||
timeout=account_timeout,
|
||||
)
|
||||
logger.success(f"[{task_type}] ({idx}/{total}) {login} — done")
|
||||
async with lock:
|
||||
self._account_results[task_id][acc_id] = {"status": "ok"}
|
||||
except asyncio.TimeoutError:
|
||||
async with lock:
|
||||
errors.append(f"{login}: timed out after {account_timeout}s")
|
||||
self._account_results[task_id][acc_id] = {"status": "error", "error": f"Timed out after {account_timeout}s"}
|
||||
logger.error(f"[{task_type}] ({idx}/{total}) {login} — timed out after {account_timeout}s")
|
||||
except Exception as exc:
|
||||
err_str = str(exc)
|
||||
async with lock:
|
||||
errors.append(f"{login}: {exc}")
|
||||
self._account_results[task_id][acc_id] = {"status": "error", "error": err_str}
|
||||
logger.error(f"[{task_type}] ({idx}/{total}) {login} — error: {exc}")
|
||||
finally:
|
||||
async with lock:
|
||||
self._active_counts[task_id] = max(0, self._active_counts.get(task_id, 0) - 1)
|
||||
|
||||
async with lock:
|
||||
completed += 1
|
||||
db = await get_db()
|
||||
await db.execute(
|
||||
"UPDATE tasks SET progress = ?, updated_at = datetime('now') WHERE id = ?",
|
||||
(completed, task_id),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
try:
|
||||
tasks = [process_one(acc, i) for i, acc in enumerate(accounts, 1)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
db = await get_db()
|
||||
status = "completed"
|
||||
result = f"Done: {completed - len(errors)} success, {len(errors)} errors"
|
||||
error_text = "; ".join(errors[:10]) if errors else None
|
||||
results_json = json.dumps(
|
||||
{str(k): v for k, v in self._account_results.get(task_id, {}).items()}
|
||||
)
|
||||
await db.execute(
|
||||
"UPDATE tasks SET status = ?, result = ?, error = ?, account_results = ?, updated_at = datetime('now') WHERE id = ?",
|
||||
(status, result, error_text, results_json, task_id),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
except asyncio.CancelledError:
|
||||
db = await get_db()
|
||||
results_json = json.dumps(
|
||||
{str(k): v for k, v in self._account_results.get(task_id, {}).items()}
|
||||
)
|
||||
await db.execute(
|
||||
"UPDATE tasks SET status = 'cancelled', account_results = ?, updated_at = datetime('now') WHERE id = ?",
|
||||
(results_json, task_id),
|
||||
)
|
||||
await db.commit()
|
||||
except Exception as exc:
|
||||
logger.error(f"Task {task_id} fatal error: {exc}")
|
||||
db = await get_db()
|
||||
results_json = json.dumps(
|
||||
{str(k): v for k, v in self._account_results.get(task_id, {}).items()}
|
||||
)
|
||||
await db.execute(
|
||||
"UPDATE tasks SET status = 'failed', error = ?, account_results = ?, updated_at = datetime('now') WHERE id = ?",
|
||||
(str(exc)[:500], results_json, task_id),
|
||||
)
|
||||
await db.commit()
|
||||
finally:
|
||||
self._running.pop(task_id, None)
|
||||
self._steps.pop(task_id, None)
|
||||
self._active_counts.pop(task_id, None)
|
||||
self._account_results.pop(task_id, None)
|
||||
self._account_steps.pop(task_id, None)
|
||||
|
||||
def get_active_count(self, task_id: str) -> int:
|
||||
return self._active_counts.get(task_id, 0)
|
||||
|
||||
def get_account_results(self, task_id: str) -> dict[int, dict[str, str]]:
|
||||
return self._account_results.get(task_id, {})
|
||||
|
||||
def clear_account_results(self, task_id: str) -> None:
|
||||
self._account_results.pop(task_id, None)
|
||||
|
||||
async def set_step(self, task_id: str, step: int, total_steps: int, label: str = "", acc_id: int = 0) -> None:
|
||||
self._steps[task_id] = {"step": step, "total_steps": total_steps, "label": label}
|
||||
if acc_id:
|
||||
self._account_steps.setdefault(task_id, {})[acc_id] = {"step": step, "total": total_steps}
|
||||
|
||||
def get_step_info(self, task_id: str) -> dict[str, Any] | None:
|
||||
return self._steps.get(task_id)
|
||||
|
||||
def get_account_steps(self, task_id: str) -> dict[int, dict[str, int]]:
|
||||
return self._account_steps.get(task_id, {})
|
||||
|
||||
async def prompt_user(self, task_id: str, message: str, login: str = "") -> str:
|
||||
"""Pause task execution and ask the user for input via SSE prompt."""
|
||||
key = f"{task_id}:{login}" if login else task_id
|
||||
event = asyncio.Event()
|
||||
self._prompts[key] = {"message": message, "login": login, "event": event, "response": ""}
|
||||
await event.wait()
|
||||
response = self._prompts.pop(key, {}).get("response", "")
|
||||
return response
|
||||
|
||||
def respond(self, task_id: str, value: str, login: str = "") -> bool:
|
||||
key = f"{task_id}:{login}" if login else task_id
|
||||
prompt = self._prompts.get(key)
|
||||
if not prompt:
|
||||
return False
|
||||
prompt["response"] = value
|
||||
prompt["event"].set()
|
||||
return True
|
||||
|
||||
def get_pending_prompt(self, task_id: str) -> dict[str, str] | None:
|
||||
"""Return first pending prompt for this task (with login info)."""
|
||||
prefix = f"{task_id}:"
|
||||
for key, prompt in self._prompts.items():
|
||||
if key == task_id or key.startswith(prefix):
|
||||
return {"message": prompt["message"], "login": prompt.get("login", "")}
|
||||
return None
|
||||
|
||||
def cancel(self, task_id: str) -> bool:
|
||||
task = self._running.get(task_id)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
return True
|
||||
return False
|
||||
|
||||
def active_count(self) -> int:
|
||||
return len(self._running)
|
||||
|
||||
|
||||
task_manager = TaskManager()
|
||||
Reference in New Issue
Block a user