422 lines
14 KiB
Python
422 lines
14 KiB
Python
"""
|
|
Weibo-HotSign Sign-in Task Definitions
|
|
Celery tasks for scheduled sign-in operations
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import asyncio
|
|
import httpx
|
|
import json
|
|
import logging
|
|
import redis
|
|
from datetime import datetime
|
|
from typing import Dict, Any, Optional
|
|
|
|
from celery import current_task
|
|
from sqlalchemy import select
|
|
|
|
# Add parent directory to path for imports
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../.."))
|
|
|
|
from shared.models.base import AsyncSessionLocal
|
|
from shared.models.task import Task
|
|
from shared.models.account import Account
|
|
from shared.config import shared_settings
|
|
|
|
from ..celery_app import celery_app
|
|
from ..config import settings
|
|
|
|
# Configure logger
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Redis client for distributed locks (可选)
|
|
_redis_client = None
|
|
|
|
def get_redis_client():
|
|
"""获取 Redis 客户端，如果未启用则返回 None"""
|
|
global _redis_client
|
|
if not shared_settings.USE_REDIS:
|
|
return None
|
|
|
|
if _redis_client is None:
|
|
try:
|
|
_redis_client = redis.from_url(shared_settings.REDIS_URL, decode_responses=True)
|
|
except Exception as e:
|
|
logger.warning(f"Redis 连接失败: {e}，分布式锁将被禁用")
|
|
return None
|
|
return _redis_client
|
|
|
|
|
|
# 内存锁（当 Redis 不可用时）
|
|
_memory_locks = {}
|
|
|
|
|
|
class DistributedLock:
|
|
"""分布式锁（支持 Redis 或内存模式）"""
|
|
|
|
def __init__(self, lock_key: str, timeout: int = 300):
|
|
"""
|
|
Initialize distributed lock
|
|
|
|
Args:
|
|
lock_key: Unique key for the lock
|
|
timeout: Lock timeout in seconds (default 5 minutes)
|
|
"""
|
|
self.lock_key = f"lock:{lock_key}"
|
|
self.timeout = timeout
|
|
self.acquired = False
|
|
self.redis_client = get_redis_client()
|
|
|
|
def acquire(self) -> bool:
|
|
"""
|
|
Acquire the lock using Redis SETNX or memory dict
|
|
Returns True if lock acquired, False otherwise
|
|
"""
|
|
try:
|
|
if self.redis_client:
|
|
# 使用 Redis
|
|
result = self.redis_client.set(self.lock_key, "1", nx=True, ex=self.timeout)
|
|
self.acquired = bool(result)
|
|
else:
|
|
# 使用内存锁（本地开发）
|
|
if self.lock_key not in _memory_locks:
|
|
_memory_locks[self.lock_key] = True
|
|
self.acquired = True
|
|
else:
|
|
self.acquired = False
|
|
|
|
return self.acquired
|
|
except Exception as e:
|
|
logger.error(f"Failed to acquire lock {self.lock_key}: {e}")
|
|
return False
|
|
|
|
def release(self):
|
|
"""Release the lock"""
|
|
if self.acquired:
|
|
try:
|
|
if self.redis_client:
|
|
# 使用 Redis
|
|
self.redis_client.delete(self.lock_key)
|
|
else:
|
|
# 使用内存锁
|
|
if self.lock_key in _memory_locks:
|
|
del _memory_locks[self.lock_key]
|
|
|
|
self.acquired = False
|
|
except Exception as e:
|
|
logger.error(f"Failed to release lock {self.lock_key}: {e}")
|
|
|
|
def __enter__(self):
|
|
"""Context manager entry"""
|
|
if not self.acquire():
|
|
raise Exception(f"Failed to acquire lock: {self.lock_key}")
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Context manager exit"""
|
|
self.release()
|
|
|
|
@celery_app.task(bind=True, max_retries=3, default_retry_delay=60)
|
|
def execute_signin_task(self, task_id: str, account_id: str, cron_expression: str):
|
|
"""
|
|
Execute scheduled sign-in task for a specific account
|
|
This task is triggered by Celery Beat based on cron schedule
|
|
Uses distributed lock to prevent duplicate execution
|
|
"""
|
|
lock_key = f"signin_task:{task_id}:{account_id}"
|
|
lock = DistributedLock(lock_key, timeout=300)
|
|
|
|
# Try to acquire lock
|
|
if not lock.acquire():
|
|
logger.warning(f"⚠️ Task {task_id} for account {account_id} is already running, skipping")
|
|
return {
|
|
"status": "skipped",
|
|
"reason": "Task already running (distributed lock)",
|
|
"account_id": account_id,
|
|
"task_id": task_id
|
|
}
|
|
|
|
try:
|
|
logger.info(f"🎯 Starting sign-in task {task_id} for account {account_id}")
|
|
|
|
# Update task status
|
|
current_task.update_state(
|
|
state="PROGRESS",
|
|
meta={
|
|
"current": 10,
|
|
"total": 100,
|
|
"status": "Initializing sign-in process...",
|
|
"account_id": account_id
|
|
}
|
|
)
|
|
|
|
# Get account info from database
|
|
account_info = _get_account_from_db(account_id)
|
|
if not account_info:
|
|
raise Exception(f"Account {account_id} not found in database")
|
|
|
|
# Check if account is active
|
|
if account_info["status"] not in ["pending", "active"]:
|
|
logger.warning(f"Account {account_id} status is {account_info['status']}, skipping sign-in")
|
|
return {
|
|
"status": "skipped",
|
|
"reason": f"Account status is {account_info['status']}",
|
|
"account_id": account_id
|
|
}
|
|
|
|
# Call signin executor service
|
|
result = _call_signin_executor(account_id, task_id)
|
|
|
|
# Update task status
|
|
current_task.update_state(
|
|
state="SUCCESS",
|
|
meta={
|
|
"current": 100,
|
|
"total": 100,
|
|
"status": "Sign-in completed successfully",
|
|
"result": result,
|
|
"account_id": account_id
|
|
}
|
|
)
|
|
|
|
logger.info(f"✅ Sign-in task {task_id} completed successfully for account {account_id}")
|
|
return result
|
|
|
|
except Exception as exc:
|
|
logger.error(f"❌ Sign-in task {task_id} failed for account {account_id}: {exc}")
|
|
|
|
# Retry logic with exponential backoff
|
|
if self.request.retries < settings.MAX_RETRY_ATTEMPTS:
|
|
retry_delay = settings.RETRY_DELAY_SECONDS * (2 ** self.request.retries)
|
|
logger.info(f"🔄 Retrying task {task_id} (attempt {self.request.retries + 1}) in {retry_delay}s")
|
|
|
|
# Release lock before retry
|
|
lock.release()
|
|
|
|
raise self.retry(exc=exc, countdown=retry_delay)
|
|
|
|
# Final failure
|
|
current_task.update_state(
|
|
state="FAILURE",
|
|
meta={
|
|
"current": 100,
|
|
"total": 100,
|
|
"status": f"Task failed after {settings.MAX_RETRY_ATTEMPTS} attempts",
|
|
"error": str(exc),
|
|
"account_id": account_id
|
|
}
|
|
)
|
|
raise exc
|
|
|
|
finally:
|
|
# Always release lock
|
|
lock.release()
|
|
|
|
|
|
def _get_account_from_db(account_id: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Query account information from database (replaces mock data).
|
|
Returns account dict or None if not found.
|
|
"""
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
return loop.run_until_complete(_async_get_account(account_id))
|
|
finally:
|
|
loop.close()
|
|
|
|
|
|
async def _async_get_account(account_id: str) -> Optional[Dict[str, Any]]:
|
|
"""Async helper to query account from database"""
|
|
try:
|
|
async with AsyncSessionLocal() as session:
|
|
stmt = select(Account).where(Account.id == account_id)
|
|
result = await session.execute(stmt)
|
|
account = result.scalar_one_or_none()
|
|
|
|
if not account:
|
|
return None
|
|
|
|
return {
|
|
"id": account.id,
|
|
"user_id": account.user_id,
|
|
"weibo_user_id": account.weibo_user_id,
|
|
"remark": account.remark,
|
|
"status": account.status,
|
|
"encrypted_cookies": account.encrypted_cookies,
|
|
"iv": account.iv,
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"Error querying account {account_id}: {e}")
|
|
return None
|
|
|
|
@celery_app.task
|
|
def schedule_daily_signin():
|
|
"""
|
|
Daily sign-in task - queries database for enabled tasks
|
|
"""
|
|
logger.info("📅 Executing daily sign-in schedule")
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
return loop.run_until_complete(_async_schedule_daily_signin())
|
|
finally:
|
|
loop.close()
|
|
|
|
|
|
async def _async_schedule_daily_signin():
|
|
"""Async helper to query and schedule tasks"""
|
|
try:
|
|
async with AsyncSessionLocal() as session:
|
|
# Query all enabled tasks
|
|
stmt = (
|
|
select(Task, Account)
|
|
.join(Account, Task.account_id == Account.id)
|
|
.where(Task.is_enabled == True)
|
|
.where(Account.status.in_(["pending", "active"]))
|
|
)
|
|
result = await session.execute(stmt)
|
|
task_account_pairs = result.all()
|
|
|
|
results = []
|
|
for task, account in task_account_pairs:
|
|
try:
|
|
# Submit individual sign-in task for each account
|
|
celery_task = execute_signin_task.delay(
|
|
task_id=task.id,
|
|
account_id=account.id,
|
|
cron_expression=task.cron_expression
|
|
)
|
|
results.append({
|
|
"account_id": account.id,
|
|
"task_id": celery_task.id,
|
|
"status": "submitted"
|
|
})
|
|
except Exception as e:
|
|
logger.error(f"Failed to submit task for account {account.id}: {e}")
|
|
results.append({
|
|
"account_id": account.id,
|
|
"status": "failed",
|
|
"error": str(e)
|
|
})
|
|
|
|
return {
|
|
"scheduled_date": datetime.now().isoformat(),
|
|
"accounts_processed": len(task_account_pairs),
|
|
"results": results
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"Error in daily signin schedule: {e}")
|
|
raise
|
|
|
|
@celery_app.task
|
|
def process_pending_tasks():
|
|
"""
|
|
Process pending sign-in tasks from database
|
|
Queries database for enabled tasks and submits them for execution
|
|
"""
|
|
logger.info("🔄 Processing pending sign-in tasks from database")
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
return loop.run_until_complete(_async_process_pending_tasks())
|
|
finally:
|
|
loop.close()
|
|
|
|
|
|
async def _async_process_pending_tasks():
|
|
"""Async helper to process pending tasks"""
|
|
try:
|
|
async with AsyncSessionLocal() as session:
|
|
# Query enabled tasks that are due for execution
|
|
stmt = (
|
|
select(Task, Account)
|
|
.join(Account, Task.account_id == Account.id)
|
|
.where(Task.is_enabled == True)
|
|
.where(Account.status.in_(["pending", "active"]))
|
|
)
|
|
result = await session.execute(stmt)
|
|
task_account_pairs = result.all()
|
|
|
|
tasks_submitted = 0
|
|
tasks_skipped = 0
|
|
|
|
for task, account in task_account_pairs:
|
|
try:
|
|
# Submit task for execution
|
|
execute_signin_task.delay(
|
|
task_id=task.id,
|
|
account_id=account.id,
|
|
cron_expression=task.cron_expression
|
|
)
|
|
tasks_submitted += 1
|
|
except Exception as e:
|
|
logger.error(f"Failed to submit task {task.id}: {e}")
|
|
tasks_skipped += 1
|
|
|
|
result = {
|
|
"processed_at": datetime.now().isoformat(),
|
|
"tasks_found": len(task_account_pairs),
|
|
"tasks_submitted": tasks_submitted,
|
|
"tasks_skipped": tasks_skipped,
|
|
"status": "completed"
|
|
}
|
|
|
|
logger.info(f"✅ Processed pending tasks: {result}")
|
|
return result
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ Failed to process pending tasks: {e}")
|
|
raise
|
|
|
|
def _call_signin_executor(account_id: str, task_id: str) -> Dict[str, Any]:
|
|
"""
|
|
Call the signin executor service to perform actual sign-in
|
|
"""
|
|
try:
|
|
signin_data = {
|
|
"task_id": task_id,
|
|
"account_id": account_id,
|
|
"timestamp": datetime.now().isoformat(),
|
|
"requested_by": "task_scheduler"
|
|
}
|
|
|
|
# Call signin executor service
|
|
with httpx.Client(timeout=30.0) as client:
|
|
response = client.post(
|
|
f"{settings.SIGNIN_EXECUTOR_URL}/api/v1/signin/execute",
|
|
json=signin_data
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
result = response.json()
|
|
logger.info(f"Sign-in executor response: {result}")
|
|
return result
|
|
else:
|
|
raise Exception(f"Sign-in executor returned error: {response.status_code} - {response.text}")
|
|
|
|
except httpx.RequestError as e:
|
|
logger.error(f"Network error calling signin executor: {e}")
|
|
raise Exception(f"Failed to connect to signin executor: {e}")
|
|
except Exception as e:
|
|
logger.error(f"Error calling signin executor: {e}")
|
|
raise
|
|
|
|
# Periodic task definitions for Celery Beat
|
|
celery_app.conf.beat_schedule = {
|
|
"daily-signin-at-8am": {
|
|
"task": "app.tasks.signin_tasks.schedule_daily_signin",
|
|
"schedule": {
|
|
"hour": 8,
|
|
"minute": 0,
|
|
},
|
|
},
|
|
"process-pending-every-15-minutes": {
|
|
"task": "app.tasks.signin_tasks.process_pending_tasks",
|
|
"schedule": 900.0, # Every 15 minutes
|
|
},
|
|
}
|