#!/usr/bin/env python3 """ Smart API queue worker. - Claims jobs with FOR UPDATE SKIP LOCKED style (status flip) - Priority-aware - Concurrent limited - Exponential backoff on failures - Extracts tags from comments, counts replies """ from __future__ import annotations import asyncio import json import os import re import signal import sys from datetime import datetime, timedelta from typing import Any, Optional import aiomysql from dotenv import load_dotenv from loguru import logger from instagram_client import InstagramClient load_dotenv() WORKER_ID = os.getenv("WORKER_ID", "worker-1") POLL_INTERVAL = float(os.getenv("POLL_INTERVAL", "0.8")) LOCK_TIMEOUT = int(os.getenv("LOCK_TIMEOUT", "60")) MAX_CONCURRENT = int(os.getenv("MAX_CONCURRENT", "3")) DB_CFG = { "host": os.getenv("DB_HOST", "127.0.0.1"), "port": int(os.getenv("DB_PORT", "3306")), "user": os.getenv("DB_USER", "root"), "password": os.getenv("DB_PASS", ""), "db": os.getenv("DB_NAME", "konkurs_bot"), "charset": "utf8mb4", "autocommit": False, } TAG_RE = re.compile(r"@([A-Za-z0-9._]+)") running = True sem: Optional[asyncio.Semaphore] = None ig: Optional[InstagramClient] = None def stop(*_): global running running = False logger.info("Shutdown signal received") async def get_pool(): return await aiomysql.create_pool(minsize=1, maxsize=MAX_CONCURRENT + 2, **DB_CFG) async def claim_job(pool) -> Optional[dict]: """Atomically claim next pending job by priority.""" async with pool.acquire() as conn: async with conn.cursor(aiomysql.DictCursor) as cur: await conn.begin() try: # Recover stale locks await cur.execute( """ UPDATE api_queue SET status='pending', locked_at=NULL, locked_by=NULL WHERE status='processing' AND locked_at < NOW() - INTERVAL %s SECOND """, (LOCK_TIMEOUT,), ) await cur.execute( """ SELECT id FROM api_queue WHERE status='pending' AND attempts < max_attempts ORDER BY priority ASC, created_at ASC LIMIT 1 FOR UPDATE """ ) row = await cur.fetchone() if not row: await conn.commit() return None job_id = row["id"] await cur.execute( """ UPDATE api_queue SET status='processing', locked_at=NOW(), locked_by=%s, attempts=attempts+1 WHERE id=%s """, (WORKER_ID, job_id), ) await cur.execute("SELECT * FROM api_queue WHERE id=%s", (job_id,)) job = await cur.fetchone() await conn.commit() return job except Exception: await conn.rollback() raise async def finish_job(pool, job_id: int, result: Any = None, error: str | None = None): async with pool.acquire() as conn: async with conn.cursor() as cur: if error: await cur.execute( """ UPDATE api_queue SET status=IF(attempts>=max_attempts,'failed','pending'), error_message=%s, result=NULL, locked_at=NULL, locked_by=NULL, updated_at=NOW() WHERE id=%s """, (error[:500], job_id), ) else: payload = json.dumps(result, ensure_ascii=False, default=str) await cur.execute( """ UPDATE api_queue SET status='done', result=%s, error_message=NULL, locked_at=NULL, locked_by=NULL, updated_at=NOW() WHERE id=%s """, (payload, job_id), ) await conn.commit() async def process_job(job: dict) -> dict[str, Any]: assert ig is not None task = job["task_type"] payload = job["payload"] if isinstance(payload, (bytes, str)): payload = json.loads(payload) if task == "resolve_user": return await ig.resolve_user(payload["username"]) if task == "check_follow": data = await ig.check_follow( payload.get("store_username") or payload.get("target_username", ""), str(payload["instagram_id"]), ) return {"is_following": bool(data.get("is_following"))} if task == "check_like": data = await ig.check_like(payload["instagram_username"], payload["media_url"]) check = data.get("check_result") or {} return {"did_like": bool(check.get("did_like"))} if task == "check_comment": data = await ig.check_comment(payload["instagram_username"], payload["media_url"]) results = data.get("results") or [] if not results: return {"did_comment": False} item = results[0] did = bool(item.get("did_comment")) last = item.get("last_comment") or {} text = last.get("text") or "" comment_id = last.get("comment_id") tags = TAG_RE.findall(text) # unique tags excluding self self_u = (payload.get("instagram_username") or "").lower() tags = list(dict.fromkeys(t.lower() for t in tags if t.lower() != self_u)) keyword = (payload.get("comment_keyword") or "").strip() kw_ok = True if keyword and did: kw_ok = keyword.lower() in text.lower() min_tags = int(payload.get("min_tags") or 0) tags_ok = len(tags) >= min_tags if min_tags else True return { "did_comment": did and kw_ok and tags_ok, "comment_text": text, "comment_id": comment_id, "tagged_usernames": tags, "keyword_ok": kw_ok, "tags_ok": tags_ok, } if task == "check_replies": media_id = payload.get("media_id") comment_id = payload.get("comment_id") if not media_id or not comment_id: return {"reply_count": 0, "error": "missing media_id or comment_id"} data = await ig.comment_replies(str(media_id), str(comment_id), max_pages=3) replies = data.get("replies") or [] # Count unique users (not self) self_u = (payload.get("instagram_username") or "").lower() users = set() for r in replies: u = ((r.get("user") or {}).get("username") or "").lower() if u and u != self_u: users.add(u) meta = data.get("meta") or {} total = meta.get("reply_count") or len(replies) return { "reply_count": max(len(users), int(total) if total else 0), "unique_repliers": list(users), } if task == "media_info": return await ig.media_info(payload["media_url"]) raise RuntimeError(f"Unknown task: {task}") async def handle(pool, job: dict): job_id = job["id"] try: logger.info(f"[{WORKER_ID}] job#{job_id} {job['task_type']}") result = await process_job(job) await finish_job(pool, job_id, result=result) logger.info(f"[{WORKER_ID}] job#{job_id} done") except Exception as e: logger.exception(f"[{WORKER_ID}] job#{job_id} error: {e}") await finish_job(pool, job_id, error=str(e)) async def worker_loop(pool): global sem sem = asyncio.Semaphore(MAX_CONCURRENT) tasks: set[asyncio.Task] = set() while running: # Clean finished tasks done = {t for t in tasks if t.done()} for t in done: try: t.result() except Exception as e: logger.error(f"Task error: {e}") tasks -= done if len(tasks) >= MAX_CONCURRENT: await asyncio.sleep(0.3) continue job = await claim_job(pool) if not job: await asyncio.sleep(POLL_INTERVAL) continue async def _run(j): async with sem: await handle(pool, j) t = asyncio.create_task(_run(job)) tasks.add(t) if tasks: await asyncio.gather(*tasks, return_exceptions=True) async def main(): global ig logger.remove() logger.add(sys.stderr, level="INFO") logger.add("../storage/logs/worker.log", rotation="10 MB", retention="7 days", level="DEBUG") signal.signal(signal.SIGINT, stop) signal.signal(signal.SIGTERM, stop) ig = InstagramClient() pool = await get_pool() logger.info(f"Worker {WORKER_ID} started (max_concurrent={MAX_CONCURRENT})") try: await worker_loop(pool) finally: await ig.close() pool.close() await pool.wait_closed() logger.info("Worker stopped") if __name__ == "__main__": asyncio.run(main())