|
| 1 | +"""Shared platform monitor dispatch — extracted from CharlesSchwabPlatform/FirstradePlatform/LongBridgePlatform (byte-identical).""" |
| 2 | + |
| 3 | +"""Dispatch shared monitor checks to platform Cloud Run services.""" |
| 4 | + |
| 5 | +from __future__ import annotations |
| 6 | + |
| 7 | +import datetime as dt |
| 8 | +import json |
| 9 | +import os |
| 10 | +from concurrent.futures import ThreadPoolExecutor, as_completed |
| 11 | +from dataclasses import dataclass |
| 12 | +from typing import Any, Callable, Iterable, Mapping, Sequence |
| 13 | +from zoneinfo import ZoneInfo |
| 14 | + |
| 15 | +import requests |
| 16 | +from ..cloud import get_deployment_context |
| 17 | + |
| 18 | + |
| 19 | +MONITOR_TARGET_ENV_NAMES = ( |
| 20 | + "MONITOR_DISPATCH_TARGETS_JSON", |
| 21 | + "LONGBRIDGE_MONITOR_DISPATCH_TARGETS_JSON", |
| 22 | + "SCHWAB_MONITOR_DISPATCH_TARGETS_JSON", |
| 23 | + "FIRSTRADE_MONITOR_DISPATCH_TARGETS_JSON", |
| 24 | +) |
| 25 | +DEFAULT_LOOKBACK_MINUTES = 4 |
| 26 | +DEFAULT_TIMEOUT_SECONDS = 120 |
| 27 | +DEFAULT_MAX_WORKERS = 4 |
| 28 | + |
| 29 | + |
| 30 | +@dataclass(frozen=True) |
| 31 | +class MonitorWindow: |
| 32 | + name: str |
| 33 | + path: str |
| 34 | + scheduler_key: str |
| 35 | + |
| 36 | + |
| 37 | +MONITOR_WINDOWS = ( |
| 38 | + MonitorWindow("probe", "/probe", "probe_time"), |
| 39 | + MonitorWindow("precheck", "/dry-run", "precheck_time"), |
| 40 | +) |
| 41 | + |
| 42 | + |
| 43 | +def load_monitor_targets(env: Mapping[str, str] | None = None) -> list[dict[str, Any]]: |
| 44 | + env = env or os.environ |
| 45 | + raw = "" |
| 46 | + for name in MONITOR_TARGET_ENV_NAMES: |
| 47 | + raw = str(env.get(name) or "").strip() |
| 48 | + if raw: |
| 49 | + break |
| 50 | + if not raw: |
| 51 | + return [] |
| 52 | + payload = json.loads(raw) |
| 53 | + if isinstance(payload, dict): |
| 54 | + targets = payload.get("targets") |
| 55 | + else: |
| 56 | + targets = payload |
| 57 | + if not isinstance(targets, list): |
| 58 | + raise ValueError("monitor dispatch targets must be a JSON array or an object with targets") |
| 59 | + return [dict(target) for target in targets if isinstance(target, Mapping)] |
| 60 | + |
| 61 | + |
| 62 | +def dispatch_due_monitors( |
| 63 | + targets: Sequence[Mapping[str, Any]], |
| 64 | + *, |
| 65 | + now: dt.datetime | None = None, |
| 66 | + lookback_minutes: int = DEFAULT_LOOKBACK_MINUTES, |
| 67 | + timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS, |
| 68 | + max_workers: int = DEFAULT_MAX_WORKERS, |
| 69 | + token_fetcher: Callable[[str], str] | None = None, |
| 70 | + post_fn: Callable[..., Any] | None = None, |
| 71 | +) -> dict[str, Any]: |
| 72 | + now_utc = _as_utc(now or dt.datetime.now(dt.timezone.utc)) |
| 73 | + due_dispatches = list(_iter_due_dispatches(targets, now_utc=now_utc, lookback_minutes=lookback_minutes)) |
| 74 | + token_fetcher = token_fetcher or _fetch_id_token |
| 75 | + post_fn = post_fn or requests.post |
| 76 | + if not due_dispatches: |
| 77 | + return { |
| 78 | + "ok": True, |
| 79 | + "total_targets": len(targets), |
| 80 | + "dispatches_due": 0, |
| 81 | + "results": [], |
| 82 | + } |
| 83 | + |
| 84 | + workers = max(1, min(int(max_workers or 1), len(due_dispatches))) |
| 85 | + results: list[dict[str, Any]] = [] |
| 86 | + with ThreadPoolExecutor(max_workers=workers) as executor: |
| 87 | + futures = [ |
| 88 | + executor.submit( |
| 89 | + _send_dispatch, |
| 90 | + dispatch, |
| 91 | + timeout_seconds=timeout_seconds, |
| 92 | + token_fetcher=token_fetcher, |
| 93 | + post_fn=post_fn, |
| 94 | + ) |
| 95 | + for dispatch in due_dispatches |
| 96 | + ] |
| 97 | + for future in as_completed(futures): |
| 98 | + results.append(future.result()) |
| 99 | + results.sort(key=lambda item: (str(item.get("service_name") or ""), str(item.get("window") or ""))) |
| 100 | + return { |
| 101 | + "ok": all(bool(result.get("ok")) for result in results), |
| 102 | + "total_targets": len(targets), |
| 103 | + "dispatches_due": len(due_dispatches), |
| 104 | + "results": results, |
| 105 | + } |
| 106 | + |
| 107 | + |
| 108 | +def _iter_due_dispatches( |
| 109 | + targets: Sequence[Mapping[str, Any]], |
| 110 | + *, |
| 111 | + now_utc: dt.datetime, |
| 112 | + lookback_minutes: int, |
| 113 | +) -> Iterable[dict[str, Any]]: |
| 114 | + for target in targets: |
| 115 | + if not _target_enabled(target): |
| 116 | + continue |
| 117 | + service_url = str(target.get("service_url") or "").rstrip("/") |
| 118 | + if not service_url: |
| 119 | + continue |
| 120 | + scheduler = target.get("scheduler") if isinstance(target.get("scheduler"), Mapping) else {} |
| 121 | + timezone = _target_timezone(scheduler) |
| 122 | + local_now = now_utc.astimezone(timezone) |
| 123 | + for window in MONITOR_WINDOWS: |
| 124 | + schedule = str(scheduler.get(window.scheduler_key) or "").strip() |
| 125 | + if not schedule: |
| 126 | + continue |
| 127 | + if _cron_due_within_window(schedule, local_now=local_now, lookback_minutes=lookback_minutes): |
| 128 | + yield { |
| 129 | + "service_name": str(target.get("service_name") or ""), |
| 130 | + "strategy_profile": str(target.get("strategy_profile") or ""), |
| 131 | + "account_scope": str(target.get("account_scope") or target.get("account_group") or ""), |
| 132 | + "window": window.name, |
| 133 | + "url": f"{service_url}{window.path}", |
| 134 | + "audience": service_url, |
| 135 | + } |
| 136 | + |
| 137 | + |
| 138 | +def _send_dispatch( |
| 139 | + dispatch: Mapping[str, Any], |
| 140 | + *, |
| 141 | + timeout_seconds: int, |
| 142 | + token_fetcher: Callable[[str], str], |
| 143 | + post_fn: Callable[..., Any], |
| 144 | +) -> dict[str, Any]: |
| 145 | + url = str(dispatch.get("url") or "") |
| 146 | + audience = str(dispatch.get("audience") or "") |
| 147 | + base_result = { |
| 148 | + "service_name": dispatch.get("service_name"), |
| 149 | + "strategy_profile": dispatch.get("strategy_profile"), |
| 150 | + "account_scope": dispatch.get("account_scope"), |
| 151 | + "window": dispatch.get("window"), |
| 152 | + "url": url, |
| 153 | + } |
| 154 | + try: |
| 155 | + token = token_fetcher(audience) |
| 156 | + response = post_fn( |
| 157 | + url, |
| 158 | + headers={ |
| 159 | + "Authorization": f"Bearer {token}", |
| 160 | + "User-Agent": "platform-monitor-dispatcher", |
| 161 | + }, |
| 162 | + timeout=timeout_seconds, |
| 163 | + ) |
| 164 | + status_code = int(getattr(response, "status_code", 0) or 0) |
| 165 | + return { |
| 166 | + **base_result, |
| 167 | + "status_code": status_code, |
| 168 | + "ok": 200 <= status_code < 300, |
| 169 | + } |
| 170 | + except Exception as exc: |
| 171 | + return { |
| 172 | + **base_result, |
| 173 | + "status_code": 0, |
| 174 | + "ok": False, |
| 175 | + "error_type": type(exc).__name__, |
| 176 | + "error": str(exc), |
| 177 | + } |
| 178 | + |
| 179 | + |
| 180 | +def _target_enabled(target: Mapping[str, Any]) -> bool: |
| 181 | + value = target.get("runtime_target_enabled") |
| 182 | + if value is None: |
| 183 | + value = target.get("enabled") |
| 184 | + if value is None: |
| 185 | + return True |
| 186 | + if isinstance(value, bool): |
| 187 | + return value |
| 188 | + return str(value).strip().lower() not in {"0", "false", "no", "off", "disabled"} |
| 189 | + |
| 190 | + |
| 191 | +def _target_timezone(scheduler: Mapping[str, Any]) -> dt.tzinfo: |
| 192 | + try: |
| 193 | + return ZoneInfo(str(scheduler.get("timezone") or "UTC")) |
| 194 | + except Exception: |
| 195 | + return dt.timezone.utc |
| 196 | + |
| 197 | + |
| 198 | +def _cron_due_within_window(schedule: str, *, local_now: dt.datetime, lookback_minutes: int) -> bool: |
| 199 | + fields = schedule.split() |
| 200 | + if len(fields) != 5: |
| 201 | + return False |
| 202 | + lookback = max(0, int(lookback_minutes or 0)) |
| 203 | + floor_now = local_now.replace(second=0, microsecond=0) |
| 204 | + for minute_offset in range(lookback + 1): |
| 205 | + candidate = floor_now - dt.timedelta(minutes=minute_offset) |
| 206 | + if _cron_matches(fields, candidate): |
| 207 | + return True |
| 208 | + return False |
| 209 | + |
| 210 | + |
| 211 | +def _cron_matches(fields: Sequence[str], value: dt.datetime) -> bool: |
| 212 | + minute, hour, day_of_month, month, day_of_week = fields |
| 213 | + cron_weekday = (value.weekday() + 1) % 7 |
| 214 | + return ( |
| 215 | + _field_matches(minute, value.minute, 0, 59) |
| 216 | + and _field_matches(hour, value.hour, 0, 23) |
| 217 | + and _field_matches(day_of_month, value.day, 1, 31) |
| 218 | + and _field_matches(month, value.month, 1, 12) |
| 219 | + and _field_matches(day_of_week, cron_weekday, 0, 7, sunday_alias=True) |
| 220 | + ) |
| 221 | + |
| 222 | + |
| 223 | +def _field_matches(field: str, value: int, min_value: int, max_value: int, *, sunday_alias: bool = False) -> bool: |
| 224 | + for part in field.split(","): |
| 225 | + part = part.strip() |
| 226 | + if not part: |
| 227 | + continue |
| 228 | + if _part_matches(part, value, min_value, max_value, sunday_alias=sunday_alias): |
| 229 | + return True |
| 230 | + return False |
| 231 | + |
| 232 | + |
| 233 | +def _part_matches(part: str, value: int, min_value: int, max_value: int, *, sunday_alias: bool = False) -> bool: |
| 234 | + if "/" in part: |
| 235 | + base, step_text = part.split("/", 1) |
| 236 | + try: |
| 237 | + step = int(step_text) |
| 238 | + except ValueError: |
| 239 | + return False |
| 240 | + if step <= 0: |
| 241 | + return False |
| 242 | + else: |
| 243 | + base = part |
| 244 | + step = 1 |
| 245 | + if base == "*": |
| 246 | + start, end = min_value, max_value |
| 247 | + elif "-" in base: |
| 248 | + start_text, end_text = base.split("-", 1) |
| 249 | + try: |
| 250 | + start, end = int(start_text), int(end_text) |
| 251 | + except ValueError: |
| 252 | + return False |
| 253 | + else: |
| 254 | + try: |
| 255 | + start = end = int(base) |
| 256 | + except ValueError: |
| 257 | + return False |
| 258 | + if sunday_alias and value == 0 and start == end == 7: |
| 259 | + return True |
| 260 | + if value < start or value > end: |
| 261 | + return False |
| 262 | + return (value - start) % step == 0 |
| 263 | + |
| 264 | + |
| 265 | +def _as_utc(value: dt.datetime) -> dt.datetime: |
| 266 | + if value.tzinfo is None: |
| 267 | + return value.replace(tzinfo=dt.timezone.utc) |
| 268 | + return value.astimezone(dt.timezone.utc) |
| 269 | + |
| 270 | + |
| 271 | +def _fetch_id_token(audience: str) -> str: |
| 272 | + return get_deployment_context().fetch_id_token(audience) |
0 commit comments