501 lines
15 KiB
Python
501 lines
15 KiB
Python
import hashlib
|
|
import json
|
|
import os
|
|
import socket
|
|
import threading
|
|
import time
|
|
import uuid
|
|
|
|
import redis
|
|
|
|
from app.core.config import settings
|
|
from app.core.logging_config import get_named_logger
|
|
from app.celery_app import celery_app
|
|
|
|
logger = get_named_logger("task_scheduler")
|
|
|
|
# per-user queue scheduler:uq:{user_id}
|
|
USER_QUEUE_PREFIX = "scheduler:uq:"
|
|
# User Collection of Pending Messages
|
|
ACTIVE_USERS = "scheduler:active_users"
|
|
# Set of users that can dispatch (ready signal)
|
|
READY_SET = "scheduler:ready_users"
|
|
# Metadata of tasks that have been dispatched and are pending completion
|
|
PENDING_HASH = "scheduler:pending_tasks"
|
|
# Dynamic Sharding: Instance Registry
|
|
REGISTRY_KEY = "scheduler:instances"
|
|
|
|
TASK_TIMEOUT = 7800 # Task timeout (seconds), considered lost if exceeded
|
|
HEARTBEAT_INTERVAL = 10 # Heartbeat interval (seconds)
|
|
INSTANCE_TTL = 30 # Instance timeout (seconds)
|
|
|
|
LUA_ATOMIC_LOCK = """
|
|
local dispatch_lock = KEYS[1]
|
|
local lock_key = KEYS[2]
|
|
local instance_id = ARGV[1]
|
|
local dispatch_ttl = tonumber(ARGV[2])
|
|
local lock_ttl = tonumber(ARGV[3])
|
|
|
|
if redis.call('SET', dispatch_lock, instance_id, 'NX', 'EX', dispatch_ttl) == false then
|
|
return 0
|
|
end
|
|
|
|
if redis.call('EXISTS', lock_key) == 1 then
|
|
redis.call('DEL', dispatch_lock)
|
|
return -1
|
|
end
|
|
|
|
redis.call('SET', lock_key, 'dispatching', 'EX', lock_ttl)
|
|
return 1
|
|
"""
|
|
|
|
LUA_SAFE_DELETE = """
|
|
if redis.call('GET', KEYS[1]) == ARGV[1] then
|
|
return redis.call('DEL', KEYS[1])
|
|
end
|
|
return 0
|
|
"""
|
|
|
|
|
|
def stable_hash(value: str) -> int:
|
|
return int.from_bytes(
|
|
hashlib.md5(value.encode("utf-8")).digest(),
|
|
"big"
|
|
)
|
|
|
|
|
|
def health_check_server(scheduler_ref):
|
|
import uvicorn
|
|
from fastapi import FastAPI
|
|
|
|
health_app = FastAPI()
|
|
|
|
@health_app.get("/")
|
|
def health():
|
|
return scheduler_ref.health()
|
|
|
|
port = int(os.environ.get("SCHEDULER_HEALTH_PORT", "8001"))
|
|
threading.Thread(
|
|
target=uvicorn.run,
|
|
kwargs={
|
|
"app": health_app,
|
|
"host": "0.0.0.0",
|
|
"port": port,
|
|
"log_config": None,
|
|
},
|
|
daemon=True,
|
|
).start()
|
|
logger.info("[Health] Server started at http://0.0.0.0:%s", port)
|
|
|
|
|
|
class RedisTaskScheduler:
|
|
def __init__(self):
|
|
self.redis = redis.Redis(
|
|
host=settings.REDIS_HOST,
|
|
port=settings.REDIS_PORT,
|
|
db=settings.REDIS_DB_CELERY_BACKEND,
|
|
password=settings.REDIS_PASSWORD,
|
|
decode_responses=True,
|
|
)
|
|
self.running = False
|
|
self.dispatched = 0
|
|
self.errors = 0
|
|
|
|
self.instance_id = f"{socket.gethostname()}-{os.getpid()}"
|
|
self._shard_index = 0
|
|
self._shard_count = 1
|
|
self._last_heartbeat = 0.0
|
|
|
|
def push_task(self, task_name, user_id, params):
|
|
try:
|
|
msg_id = str(uuid.uuid4())
|
|
msg = json.dumps({
|
|
"msg_id": msg_id,
|
|
"task_name": task_name,
|
|
"user_id": user_id,
|
|
"params": json.dumps(params),
|
|
})
|
|
|
|
lock_key = f"{task_name}:{user_id}"
|
|
queue_key = f"{USER_QUEUE_PREFIX}{user_id}"
|
|
|
|
pipe = self.redis.pipeline()
|
|
pipe.rpush(queue_key, msg)
|
|
pipe.sadd(ACTIVE_USERS, user_id)
|
|
pipe.set(
|
|
f"task_tracker:{msg_id}",
|
|
json.dumps({"status": "QUEUED", "task_id": None}),
|
|
ex=86400,
|
|
)
|
|
pipe.execute()
|
|
|
|
if not self.redis.exists(lock_key):
|
|
self.redis.sadd(READY_SET, user_id)
|
|
|
|
logger.info("Task pushed: msg_id=%s task=%s user=%s", msg_id, task_name, user_id)
|
|
return msg_id
|
|
except Exception as e:
|
|
logger.error("Push task exception %s", e, exc_info=True)
|
|
raise
|
|
|
|
def get_task_status(self, msg_id: str) -> dict:
|
|
raw = self.redis.get(f"task_tracker:{msg_id}")
|
|
if raw is None:
|
|
return {"status": "NOT_FOUND"}
|
|
|
|
tracker = json.loads(raw)
|
|
status = tracker["status"]
|
|
task_id = tracker.get("task_id")
|
|
result_content = tracker.get("result") or {}
|
|
|
|
if status == "DISPATCHED" and task_id:
|
|
result_raw = self.redis.get(f"celery-task-meta-{task_id}")
|
|
if result_raw:
|
|
result_data = json.loads(result_raw)
|
|
status = result_data.get("status", status)
|
|
result_content = result_data.get("result")
|
|
|
|
return {"status": status, "task_id": task_id, "result": result_content}
|
|
|
|
def _cleanup_finished(self):
|
|
pending = self.redis.hgetall(PENDING_HASH)
|
|
if not pending:
|
|
return
|
|
|
|
now = time.time()
|
|
task_ids = list(pending.keys())
|
|
|
|
pipe = self.redis.pipeline()
|
|
for task_id in task_ids:
|
|
pipe.get(f"celery-task-meta-{task_id}")
|
|
results = pipe.execute()
|
|
|
|
cleanup_pipe = self.redis.pipeline()
|
|
has_cleanup = False
|
|
ready_user_ids = set()
|
|
|
|
for task_id, raw_result in zip(task_ids, results):
|
|
try:
|
|
meta = json.loads(pending[task_id])
|
|
lock_key = meta["lock_key"]
|
|
dispatched_at = meta.get("dispatched_at", 0)
|
|
age = now - dispatched_at
|
|
|
|
should_cleanup = False
|
|
result_data = {}
|
|
|
|
if raw_result is not None:
|
|
result_data = json.loads(raw_result)
|
|
if result_data.get("status") in ("SUCCESS", "FAILURE", "REVOKED"):
|
|
should_cleanup = True
|
|
logger.info(
|
|
"Task finished: %s state=%s", task_id,
|
|
result_data.get("status"),
|
|
)
|
|
elif age > TASK_TIMEOUT:
|
|
should_cleanup = True
|
|
logger.warning(
|
|
"Task expired or lost: %s age=%.0fs, force cleanup",
|
|
task_id, age,
|
|
)
|
|
|
|
if should_cleanup:
|
|
final_status = (
|
|
result_data.get("status", "UNKNOWN") if result_data else "EXPIRED"
|
|
)
|
|
|
|
self.redis.eval(LUA_SAFE_DELETE, 1, lock_key, task_id)
|
|
|
|
cleanup_pipe.hdel(PENDING_HASH, task_id)
|
|
|
|
tracker_msg_id = meta.get("msg_id")
|
|
if tracker_msg_id:
|
|
cleanup_pipe.set(
|
|
f"task_tracker:{tracker_msg_id}",
|
|
json.dumps({
|
|
"status": final_status,
|
|
"task_id": task_id,
|
|
"result": result_data.get("result") or {},
|
|
}),
|
|
ex=86400,
|
|
)
|
|
has_cleanup = True
|
|
|
|
parts = lock_key.split(":", 1)
|
|
if len(parts) == 2:
|
|
ready_user_ids.add(parts[1])
|
|
|
|
except Exception as e:
|
|
logger.error("Cleanup error for %s: %s", task_id, e, exc_info=True)
|
|
self.errors += 1
|
|
|
|
if has_cleanup:
|
|
cleanup_pipe.execute()
|
|
|
|
if ready_user_ids:
|
|
self.redis.sadd(READY_SET, *ready_user_ids)
|
|
|
|
def _heartbeat(self):
|
|
now = time.time()
|
|
if now - self._last_heartbeat < HEARTBEAT_INTERVAL:
|
|
return
|
|
self._last_heartbeat = now
|
|
|
|
self.redis.hset(REGISTRY_KEY, self.instance_id, str(now))
|
|
|
|
all_instances = self.redis.hgetall(REGISTRY_KEY)
|
|
|
|
alive = []
|
|
dead = []
|
|
for iid, ts in all_instances.items():
|
|
if now - float(ts) < INSTANCE_TTL:
|
|
alive.append(iid)
|
|
else:
|
|
dead.append(iid)
|
|
|
|
if dead:
|
|
pipe = self.redis.pipeline()
|
|
for iid in dead:
|
|
pipe.hdel(REGISTRY_KEY, iid)
|
|
pipe.execute()
|
|
logger.info("Cleaned dead instances: %s", dead)
|
|
|
|
alive.sort()
|
|
self._shard_count = max(len(alive), 1)
|
|
self._shard_index = (
|
|
alive.index(self.instance_id) if self.instance_id in alive else 0
|
|
)
|
|
logger.debug(
|
|
"Shard: %s/%s (instance=%s, alive=%d)",
|
|
self._shard_index, self._shard_count,
|
|
self.instance_id, len(alive),
|
|
)
|
|
|
|
def _is_mine(self, user_id: str) -> bool:
|
|
if self._shard_count <= 1:
|
|
return True
|
|
return stable_hash(user_id) % self._shard_count == self._shard_index
|
|
|
|
def _dispatch(self, msg_id, msg_data) -> bool:
|
|
user_id = msg_data["user_id"]
|
|
task_name = msg_data["task_name"]
|
|
params = json.loads(msg_data.get("params", "{}"))
|
|
|
|
lock_key = f"{task_name}:{user_id}"
|
|
dispatch_lock = f"dispatch:{msg_id}"
|
|
|
|
result = self.redis.eval(
|
|
LUA_ATOMIC_LOCK, 2,
|
|
dispatch_lock, lock_key,
|
|
self.instance_id, str(300), str(3600),
|
|
)
|
|
|
|
if result == 0:
|
|
return False
|
|
if result == -1:
|
|
return False
|
|
|
|
try:
|
|
task = celery_app.send_task(task_name, kwargs=params)
|
|
except Exception as e:
|
|
pipe = self.redis.pipeline()
|
|
pipe.delete(dispatch_lock)
|
|
pipe.delete(lock_key)
|
|
pipe.execute()
|
|
self.errors += 1
|
|
logger.error(
|
|
"send_task failed for %s:%s msg=%s: %s",
|
|
task_name, user_id, msg_id, e, exc_info=True,
|
|
)
|
|
return False
|
|
|
|
try:
|
|
pipe = self.redis.pipeline()
|
|
pipe.set(lock_key, task.id, ex=3600)
|
|
pipe.hset(PENDING_HASH, task.id, json.dumps({
|
|
"lock_key": lock_key,
|
|
"dispatched_at": time.time(),
|
|
"msg_id": msg_id,
|
|
}))
|
|
pipe.delete(dispatch_lock)
|
|
pipe.set(
|
|
f"task_tracker:{msg_id}",
|
|
json.dumps({"status": "DISPATCHED", "task_id": task.id}),
|
|
ex=86400,
|
|
)
|
|
pipe.execute()
|
|
except Exception as e:
|
|
logger.error(
|
|
"Post-dispatch state update failed for %s: %s",
|
|
task.id, e, exc_info=True,
|
|
)
|
|
self.errors += 1
|
|
|
|
self.dispatched += 1
|
|
logger.info("Task dispatched: %s (msg=%s)", task.id, msg_id)
|
|
return True
|
|
|
|
def _process_batch(self, user_ids):
|
|
if not user_ids:
|
|
return
|
|
|
|
pipe = self.redis.pipeline()
|
|
for uid in user_ids:
|
|
pipe.lindex(f"{USER_QUEUE_PREFIX}{uid}", 0)
|
|
heads = pipe.execute()
|
|
|
|
candidates = [] # (user_id, msg_dict)
|
|
empty_users = []
|
|
|
|
for uid, head in zip(user_ids, heads):
|
|
if head is None:
|
|
empty_users.append(uid)
|
|
else:
|
|
try:
|
|
candidates.append((uid, json.loads(head)))
|
|
except (json.JSONDecodeError, TypeError) as e:
|
|
logger.error("Bad message in queue for user %s: %s", uid, e)
|
|
self.redis.lpop(f"{USER_QUEUE_PREFIX}{uid}")
|
|
|
|
if empty_users:
|
|
pipe = self.redis.pipeline()
|
|
for uid in empty_users:
|
|
pipe.srem(ACTIVE_USERS, uid)
|
|
pipe.execute()
|
|
|
|
if not candidates:
|
|
return
|
|
|
|
for uid, msg in candidates:
|
|
if self._dispatch(msg["msg_id"], msg):
|
|
self.redis.lpop(f"{USER_QUEUE_PREFIX}{uid}")
|
|
|
|
def schedule_loop(self):
|
|
self._heartbeat()
|
|
self._cleanup_finished()
|
|
|
|
pipe = self.redis.pipeline()
|
|
pipe.smembers(READY_SET)
|
|
pipe.delete(READY_SET)
|
|
results = pipe.execute()
|
|
ready_users = results[0] or set()
|
|
|
|
my_users = [uid for uid in ready_users if self._is_mine(uid)]
|
|
|
|
if not my_users:
|
|
time.sleep(0.5)
|
|
return
|
|
|
|
self._process_batch(my_users)
|
|
time.sleep(0.1)
|
|
|
|
def _full_scan(self):
|
|
cursor = 0
|
|
ready_batch = []
|
|
while True:
|
|
cursor, user_ids = self.redis.sscan(
|
|
ACTIVE_USERS, cursor=cursor, count=1000,
|
|
)
|
|
if user_ids:
|
|
my_users = [uid for uid in user_ids if self._is_mine(uid)]
|
|
if my_users:
|
|
pipe = self.redis.pipeline()
|
|
for uid in my_users:
|
|
pipe.lindex(f"{USER_QUEUE_PREFIX}{uid}", 0)
|
|
heads = pipe.execute()
|
|
|
|
for uid, head in zip(my_users, heads):
|
|
if head is None:
|
|
continue
|
|
try:
|
|
msg = json.loads(head)
|
|
lock_key = f"{msg['task_name']}:{uid}"
|
|
ready_batch.append((uid, lock_key))
|
|
except (json.JSONDecodeError, TypeError):
|
|
continue
|
|
|
|
if cursor == 0:
|
|
break
|
|
|
|
if not ready_batch:
|
|
return
|
|
|
|
pipe = self.redis.pipeline()
|
|
for _, lock_key in ready_batch:
|
|
pipe.exists(lock_key)
|
|
lock_exists = pipe.execute()
|
|
|
|
ready_uids = [
|
|
uid for (uid, _), locked in zip(ready_batch, lock_exists)
|
|
if not locked
|
|
]
|
|
|
|
if ready_uids:
|
|
self.redis.sadd(READY_SET, *ready_uids)
|
|
logger.info("Full scan found %d ready users", len(ready_uids))
|
|
|
|
def run_server(self):
|
|
health_check_server(self)
|
|
self.running = True
|
|
|
|
last_full_scan = 0.0
|
|
full_scan_interval = 30.0
|
|
|
|
logger.info(
|
|
"Scheduler started: instance=%s", self.instance_id,
|
|
)
|
|
|
|
while True:
|
|
try:
|
|
self.schedule_loop()
|
|
|
|
now = time.time()
|
|
if now - last_full_scan > full_scan_interval:
|
|
self._full_scan()
|
|
last_full_scan = now
|
|
|
|
except Exception as e:
|
|
logger.error("Scheduler exception %s", e, exc_info=True)
|
|
self.errors += 1
|
|
time.sleep(5)
|
|
|
|
def health(self) -> dict:
|
|
return {
|
|
"running": self.running,
|
|
"active_users": self.redis.scard(ACTIVE_USERS),
|
|
"ready_users": self.redis.scard(READY_SET),
|
|
"pending_tasks": self.redis.hlen(PENDING_HASH),
|
|
"dispatched": self.dispatched,
|
|
"errors": self.errors,
|
|
"shard": f"{self._shard_index}/{self._shard_count}",
|
|
"instance": self.instance_id,
|
|
}
|
|
|
|
def shutdown(self):
|
|
logger.info("Scheduler shutting down: instance=%s", self.instance_id)
|
|
self.running = False
|
|
try:
|
|
self.redis.hdel(REGISTRY_KEY, self.instance_id)
|
|
except Exception as e:
|
|
logger.error("Shutdown cleanup error: %s", e)
|
|
|
|
|
|
scheduler: RedisTaskScheduler | None = None
|
|
if scheduler is None:
|
|
scheduler = RedisTaskScheduler()
|
|
|
|
if __name__ == "__main__":
|
|
import signal
|
|
import sys
|
|
|
|
|
|
def _signal_handler(signum, frame):
|
|
scheduler.shutdown()
|
|
sys.exit(0)
|
|
|
|
|
|
signal.signal(signal.SIGTERM, _signal_handler)
|
|
signal.signal(signal.SIGINT, _signal_handler)
|
|
|
|
scheduler.run_server()
|