coordinate/sentinel.py

import os
import time
import json
import random
import threading
from typing import List, Dict, Callable

from vs_types import GroupConfig, VSPlgConfig, Task, QueryWaitMode
from vs_plg_factory import VSPlgFactory 
from toolkit.thread_pool import ThreadPool 
from toolkit.vs_cloud_api import VSCloudApi
from toolkit.backoff import ExponentialBackoff
from utils.safe_redis_cli import SafeRedisClient

class SentinelGCO:
    def __init__(self, cfg: GroupConfig, redis_conf: Dict, logger: Callable[[str], None] = None):
        self.m_cfg = cfg
        self.m_factory = VSPlgFactory()
        self.m_logger = logger
        self.m_tasks: List[Task] = []
        self.m_lock = threading.RLock()
        self.m_stop_event = threading.Event()
        
        self.redis_client = SafeRedisClient(redis_conf, self.m_logger)
        self.m_pending_builtin = 0
        
        self.group_backoff = ExponentialBackoff(base_delay=60.0, max_delay=3600.0, factor=2.0)
        self.m_last_spawn_time = 0.0
        self.m_spawn_interval = 120
        self.m_last_group_query_time = 0.0

    def _log(self, message):
        if self.m_logger:
            self.m_logger(f'[SENTINEL] [{self.m_cfg.identifier}] {message}')
        else:
            print(f'[SENTINEL] [{self.m_cfg.identifier}] {message}')
            
    def _get_average_interval(self) -> float:
        """计算当前组平均的查询间隔（秒）"""
        mode = self.m_cfg.query_wait.mode
        if mode == QueryWaitMode.Loop:
            return 1.0
        elif mode == QueryWaitMode.Fixed:
            return float(self.m_cfg.query_wait.fixed_wait)
        elif mode == QueryWaitMode.Random:
            return (self.m_cfg.query_wait.random_min + self.m_cfg.query_wait.random_max) / 2.0
        return 30.0
    
    def update_config(self, new_cfg: GroupConfig):
        """
        动态更新配置
        """
        with self.m_lock:
            if self.m_cfg.enable and not new_cfg.enable:
                self._log("Config dynamically updated: Group DISABLED. Will stop creating new tasks.")
            elif not self.m_cfg.enable and new_cfg.enable:
                self._log("Config dynamically updated: Group ENABLED.")
            else:
                self._log("Config dynamically updated: Parameters refreshed.")
                
            self.m_cfg = new_cfg

    def start(self):
        if not self.m_cfg.enable:
            return
        self._log("Starting Sentinel...")
        plugin_name = self.m_cfg.plugin_config.plugin_name
        class_name = "".join(part.title() for part in plugin_name.split('_'))
        plugin_path = os.path.join(self.m_cfg.plugin_config.lib_path, self.m_cfg.plugin_config.plugin_bin)
        self.m_factory.register_plugin(plugin_name, plugin_path, class_name)

        threading.Thread(target=self._monitor_loop, daemon=True, name="Sentinel-Monitor").start()
        threading.Thread(target=self._creator_loop, daemon=True, name="Sentinel-Creator").start()

    def stop(self):
        self._log("Stopping Sentinel...")
        self.m_stop_event.set()
        with self.m_lock:
            tasks_to_cleanup = list(self.m_tasks)
            self.m_tasks.clear()
        for task in tasks_to_cleanup:
            self._cleanup_task(task, "sentinel stopped")

    def _cleanup_task(self, task: Task, reason: str):
        try:
            if task and task.instance:
                self._log(f"Cleaning up sentinel instance. reason={reason}")
                task.instance.cleanup()
        except Exception as e:
            self._log(f"Cleanup failed. reason={reason}, error={e}")

    def _remove_task(self, task: Task, reason: str):
        removed = False
        with self.m_lock:
            if task in self.m_tasks:
                self.m_tasks.remove(task)
                removed = True
        if removed:
            self._cleanup_task(task, reason)

    def _get_redis_key(self, routing_key: str) -> str:
        return f"vs:signal:{routing_key}"

    def _monitor_loop(self):
        self._log("Monitor loop started.")
        
        self.m_last_group_query_time = 0.0 
        while not self.m_stop_event.is_set():
            try:
                time.sleep(1)
                now = time.time()
                
                with self.m_lock:
                    tasks_to_check = list(self.m_tasks)

                active_tasks = []
                dead_tasks = []
                for t in tasks_to_check:
                    if t.is_querying:
                        active_tasks.append(t)
                        continue
                    
                    if t.instance.health_check():
                        active_tasks.append(t)
                    else:
                        dead_tasks.append(t)
          
                if dead_tasks:
                    with self.m_lock:
                        current_tasks = list(self.m_tasks)
                        self.m_tasks = [t for t in self.m_tasks if t in active_tasks]
                    for t in dead_tasks:
                        if t in current_tasks:
                            self._cleanup_task(t, "health check failed")
                else:
                    with self.m_lock:
                        self.m_tasks = [t for t in self.m_tasks if t in active_tasks]
                
                if not active_tasks:
                    continue

                avg_interval = self._get_average_interval()
                global_gap = max(1.0, avg_interval / len(active_tasks)) 
                
                active_tasks.sort(key=lambda x: x.next_run)
                for task in active_tasks:
                    if now < task.next_run:
                        continue
                    
                    if task.is_querying:
                        continue
                    
                    if now - self.m_last_group_query_time < global_gap:
                        break

                    apt_types = self.m_cfg.appointment_types
                    if not apt_types:
                        continue
                    weights = [float(item.weight) for item in apt_types]
                    apt_type = random.choices(apt_types, weights=weights, k=1)[0]
                    
                    interval = 30
                    mode = task.qw_cfg.mode
                    if mode == QueryWaitMode.Loop:
                        interval = 1
                    elif mode == QueryWaitMode.Fixed:
                        interval = task.qw_cfg.fixed_wait
                    elif mode == QueryWaitMode.Random:
                        interval = random.randint(task.qw_cfg.random_min, task.qw_cfg.random_max)

                    task.is_querying = True 
                    self.m_last_group_query_time = now

                    def _query_job(current_task=task, a_type=apt_type, wait_gap=interval):
                        try:
                            VSCloudApi.Instance().slot_refresh_start(a_type.routing_key, country=a_type.country, city=a_type.city, visa_type=a_type.visa_type)
                            result = current_task.instance.query(a_type)
                            result.apt_type = a_type
                            if result.success:
                                ttl = self.m_cfg.sentinel.signal_ttl
                                self._log(f"🔥 SLOT FOUND! Writing signal to Redis (TTL: {ttl}s)")
                                payload = {
                                    "group_id": self.m_cfg.identifier,
                                    "apt_type": a_type.model_dump(),
                                    "query_result": result.to_snapshot_payload(),
                                    "timestamp": time.time()
                                }
                                redis_key = self._get_redis_key(a_type.routing_key)
                                self.redis_client.setex(redis_key, ttl, json.dumps(payload))
                                payload["query_result"]["website"] = self.m_cfg.website
                                VSCloudApi.Instance().slot_snapshot_report(payload["query_result"])
                            VSCloudApi.Instance().slot_refresh_success(a_type.routing_key)
                        except Exception as e:
                            self._log(f"Query exception: {e}")
                            VSCloudApi.Instance().slot_refresh_fail(a_type.routing_key, error=str(e))
                        finally:
                            current_task.next_run = time.time() + wait_gap
                            current_task.is_querying = False
                    ThreadPool.getInstance().enqueue(_query_job)
                    break

            except Exception as e:
                self._log(f"Monitor loop error: {e}")

    def _creator_loop(self):
        self._log("Creator loop started.")
        group_cd_key = f"vs:group:cooldown:{self.m_cfg.identifier}"
        
        while not self.m_stop_event.is_set():
            try:
                time.sleep(1)
                if self.redis_client.exists(group_cd_key):
                    continue
                with self.m_lock:
                    current = len(self.m_tasks)
                    pending = self.m_pending_builtin
                    target = self.m_cfg.sentinel.target_instances
                
                if (current + pending) < target:
                    now = time.time()
                    if now - self.m_last_spawn_time >= self.m_spawn_interval:
                        with self.m_lock:
                            self.m_last_spawn_time = now
                        self._log(f"Staggered spawn triggered. Next spawn in {self.m_spawn_interval:.1f}s")
                        self._spawn_sentinel_worker()
            except Exception as e:
                self._log(f'Creator loop exception:{e}')

    def _spawn_sentinel_worker(self):
        with self.m_lock:
            self.m_pending_builtin += 1
            
        def _job():
            instance = None
            success = False
            try:
                plg_cfg = VSPlgConfig()
                plg_cfg.debug = self.m_cfg.debug
                plg_cfg.free_config = self.m_cfg.free_config
                plg_cfg.session_max_life = self.m_cfg.session_max_life
                
                if not self.m_cfg.need_account:
                    plg_cfg.account.id = 0
                    plg_cfg.account.username = "Guest"
                else:
                    acc = VSCloudApi.Instance().get_next_account(self.m_cfg.sentinel.account_pool_id, self.m_cfg.sentinel.account_cd)
                    plg_cfg.account = type(plg_cfg.account)(**acc)
                
                if self.m_cfg.need_proxy:
                    proxy = VSCloudApi.Instance().get_next_proxy(self.m_cfg.proxy_pool, self.m_cfg.proxy_cd)
                    plg_cfg.proxy = type(plg_cfg.proxy)(**proxy)

                instance = self.m_factory.create(self.m_cfg.identifier, self.m_cfg.plugin_config.plugin_name)
                instance.set_log(self.m_logger)
                instance.set_config(plg_cfg)
                instance.create_session()
                
                with self.m_lock:
                    self.m_tasks.append(
                        Task(instance=instance,qw_cfg=self.m_cfg.query_wait,next_run=time.time(), book_allowed=False))
                
                group_fail_key = f"vs:group:failures:{self.m_cfg.identifier}"
                self.redis_client.delete(group_fail_key)
                
                success = True
                self._log(f"+++ Sentinel spawned: {plg_cfg.account.username}")

            except Exception as e:
                err_str = str(e)
                resource_not_found_indicators = [
                    "40401" in err_str,
                    "Account not found" in err_str,
                    "Proxy not found" in err_str,
                ]
                if any(resource_not_found_indicators):
                    return
                
                self._log(f"Spawn failed: {e}")
                
                rate_limited_indicators = [
                    "42901" in err_str,
                    "Rate limited" in err_str
                ]
                if any(rate_limited_indicators):
                    group_fail_key = f"vs:group:failures:{self.m_cfg.identifier}"
                    group_cd_key = f"vs:group:cooldown:{self.m_cfg.identifier}"
                    
                    g_fails = self.redis_client.incr(group_fail_key)
                    g_cd = self.group_backoff.calculate(g_fails)
                    self.redis_client.set(group_cd_key, "1", ex=int(g_cd))
                    self._log(f"📉 [Rate Limited] Sentinel Spawn failed {g_fails} times. Global Backoff: {g_cd:.1f}s.")
                    
            finally:
                if not success and instance is not None:
                    instance.cleanup()
                with self.m_lock:
                    self.m_pending_builtin = max(0, self.m_pending_builtin - 1)

        ThreadPool.getInstance().enqueue(_job)