"""Redis Stream-backed realtime observation event store.""" from __future__ import annotations import json import os import socket import threading import time from typing import Any, Callable, Dict, Iterable, List, Optional, Set from loguru import logger from web.realtime_event_store import MAX_REPLAY_LIMIT, TIME_CONTRACT_KEYS from web.realtime_patch_schema import EVENT_TYPE DEFAULT_STREAM_KEY = "stream:city_observation" DEFAULT_COUNTER_KEY = "counter:city_observation_revision" DEFAULT_MAXLEN = 50000 DEFAULT_SOCKET_TIMEOUT_SECONDS = 15.0 DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS = 5.0 DEFAULT_REPLAY_CHUNK_SIZE = 512 DEFAULT_REPLAY_SCAN_MAX = 5000 APPEND_EVENT_SCRIPT = """ local revision = redis.call('INCR', KEYS[2]) local stream_id = redis.call( 'XADD', KEYS[1], 'MAXLEN', '~', ARGV[1], '*', 'revision', revision, 'type', ARGV[2], 'schema_type', ARGV[3], 'schema_version', ARGV[4], 'city', ARGV[5], 'source', ARGV[6], 'obs_time', ARGV[7], 'payload_json', ARGV[8], 'created_at_ms', ARGV[9], 'ts', ARGV[10], 'producer_id', ARGV[11] ) return {revision, stream_id} """ def _decode(value: Any) -> str: if isinstance(value, bytes): return value.decode("utf-8") return str(value or "") def _normalize_city_set(cities: Optional[Iterable[str]]) -> Set[str]: return {str(city or "").strip().lower() for city in (cities or set()) if str(city or "").strip()} def _time_contract_from_payload(payload: Dict[str, Any]) -> Dict[str, Any]: return {key: payload[key] for key in TIME_CONTRACT_KEYS if key in payload} def _int_or_zero(value: Any) -> int: try: return int(value) except (TypeError, ValueError): return 0 def _float_env(name: str, default: float) -> float: try: return float(os.getenv(name) or default) except (TypeError, ValueError): return default def _int_env(name: str, default: int, *, min_value: int, max_value: int) -> int: try: value = int(os.getenv(name) or default) except (TypeError, ValueError): value = default return max(min_value, min(max_value, int(value))) class RedisRealtimeEventStore: """Persist replayable observation patch events in a Redis Stream.""" uses_external_live_fanout = True def __init__( self, *, redis_url: Optional[str] = None, redis_client: Any = None, stream_key: Optional[str] = None, counter_key: Optional[str] = None, maxlen: Optional[int] = None, producer_id: Optional[str] = None, ) -> None: self.stream_key = stream_key or os.getenv("POLYWEATHER_REDIS_STREAM_KEY") or DEFAULT_STREAM_KEY self.counter_key = counter_key or os.getenv("POLYWEATHER_REDIS_COUNTER_KEY") or DEFAULT_COUNTER_KEY self.maxlen = max(1, int(maxlen or os.getenv("POLYWEATHER_REDIS_STREAM_MAXLEN") or DEFAULT_MAXLEN)) self.replay_chunk_size = _int_env( "POLYWEATHER_REDIS_REPLAY_CHUNK_SIZE", DEFAULT_REPLAY_CHUNK_SIZE, min_value=50, max_value=2000, ) self.replay_scan_max = _int_env( "POLYWEATHER_REDIS_REPLAY_SCAN_MAX", DEFAULT_REPLAY_SCAN_MAX, min_value=500, max_value=50000, ) self.producer_id = producer_id or os.getenv("POLYWEATHER_INSTANCE_ID") or socket.gethostname() self._client = redis_client or self._build_client(redis_url) self._subscriber_lock = threading.Lock() self._subscriber_thread: Optional[threading.Thread] = None self._subscriber_stop: Optional[threading.Event] = None @staticmethod def _build_client(redis_url: Optional[str]) -> Any: try: import redis # type: ignore except ImportError as exc: raise RuntimeError("redis package is required for Redis realtime event store") from exc url = redis_url or os.getenv("POLYWEATHER_REDIS_URL") or "redis://127.0.0.1:6379/0" client = redis.Redis.from_url( url, socket_timeout=_float_env("POLYWEATHER_REDIS_SOCKET_TIMEOUT_SECONDS", DEFAULT_SOCKET_TIMEOUT_SECONDS), socket_connect_timeout=_float_env( "POLYWEATHER_REDIS_SOCKET_CONNECT_TIMEOUT_SECONDS", DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS, ), health_check_interval=30, ) client.ping() return client def append_event(self, event: Dict[str, Any]) -> Dict[str, Any]: if event.get("type") != EVENT_TYPE: raise ValueError("unsupported realtime event type") payload = event.get("payload") if not isinstance(payload, dict): raise ValueError("event payload must be an object") created_at_ms = int(time.time() * 1000) ts = int(event.get("ts") or created_at_ms) payload_json = json.dumps(payload, ensure_ascii=False, separators=(",", ":")) result = self._client.eval( APPEND_EVENT_SCRIPT, 2, self.stream_key, self.counter_key, self.maxlen, str(event["type"]), str(event["schema_type"]), int(event["schema_version"]), str(event["city"]), str(event["source"]), str(event.get("obs_time") or ""), payload_json, created_at_ms, ts, self.producer_id, ) revision = int(_decode(result[0] if isinstance(result, (list, tuple)) else result)) return { "type": event["type"], "revision": revision, "city": str(event["city"]), "source": str(event["source"]), "obs_time": event.get("obs_time"), **_time_contract_from_payload(payload), "ts": ts, "payload": payload, } def latest_revision(self) -> int: value = self._client.get(self.counter_key) revision = _int_or_zero(_decode(value)) if revision: return revision rows = self._client.xrevrange(self.stream_key, max="+", min="-", count=1) events = self._rows_to_events(rows) return max((event["revision"] for event in events), default=0) def status(self) -> Dict[str, Any]: out: Dict[str, Any] = { "store": "redis", "redis_connected": False, "stream_key": self.stream_key, "latest_revision": 0, "stream_len": None, "oldest_revision": None, "subscriber_connected": bool( self._subscriber_thread and self._subscriber_thread.is_alive() ), } try: ping = getattr(self._client, "ping", None) if callable(ping): ping() out["redis_connected"] = True out["latest_revision"] = self.latest_revision() xlen = getattr(self._client, "xlen", None) if callable(xlen): out["stream_len"] = int(xlen(self.stream_key)) out["oldest_revision"] = self._oldest_revision() except Exception as exc: out["error"] = str(exc) return out def replay_events( self, *, cities: Optional[Set[str]], since_revision: int, limit: int, ) -> List[Dict[str, Any]]: city_set = _normalize_city_set(cities) since = max(0, int(since_revision or 0)) bounded_limit = max(1, min(MAX_REPLAY_LIMIT, int(limit or 1))) if since <= 0: return self._replay_events_forward( city_set=city_set, since=since, limit=bounded_limit, )[0] return self._replay_events_reverse( city_set=city_set, since=since, limit=bounded_limit, )[0] def replay_requires_resync( self, *, cities: Optional[Set[str]], since_revision: int, replay_count: int, limit: int, ) -> bool: city_set = _normalize_city_set(cities) since = max(0, int(since_revision or 0)) oldest_revision = self._oldest_revision() if since > 0 and oldest_revision and since < oldest_revision - 1: return True bounded_limit = max(1, int(limit or 1)) if int(replay_count or 0) < bounded_limit: return False probe_limit = min(MAX_REPLAY_LIMIT + 1, bounded_limit + 1) if since <= 0: probe_events, _, _ = self._replay_events_forward( city_set=city_set, since=since, limit=probe_limit, ) return len(probe_events) > bounded_limit probe_events, hit_boundary, scanned = self._replay_events_reverse( city_set=city_set, since=since, limit=probe_limit, ) if len(probe_events) > bounded_limit: return True return not hit_boundary and scanned >= self.replay_scan_max def start_live_subscription(self, callback: Callable[[Dict[str, Any]], None]) -> None: with self._subscriber_lock: if self._subscriber_thread and self._subscriber_thread.is_alive(): return self._subscriber_stop = threading.Event() self._subscriber_thread = threading.Thread( target=self._live_subscription_loop, args=(callback, self._subscriber_stop), name="polyweather-redis-realtime-subscriber", daemon=True, ) self._subscriber_thread.start() def stop_live_subscription(self) -> None: with self._subscriber_lock: if self._subscriber_stop: self._subscriber_stop.set() self._subscriber_thread = None self._subscriber_stop = None def _live_subscription_loop( self, callback: Callable[[Dict[str, Any]], None], stop_event: threading.Event, ) -> None: last_seen_id = "$" while not stop_event.is_set(): try: rows = self._client.xread({self.stream_key: last_seen_id}, count=100, block=5000) for _stream_name, entries in rows or []: for entry_id, fields in entries: last_seen_id = _decode(entry_id) callback(self._entry_to_event(entry_id, fields)) except Exception as exc: logger.warning(f"Redis realtime subscriber disconnected: {exc}") stop_event.wait(2.0) def _all_events(self) -> List[Dict[str, Any]]: rows = self._client.xrange(self.stream_key, min="-", max="+") return self._rows_to_events(rows) def _oldest_revision(self) -> Optional[int]: rows = self._client.xrange(self.stream_key, min="-", max="+", count=1) events = self._rows_to_events(rows) if not events: return None return int(events[0].get("revision") or 0) or None def _rows_to_events(self, rows: Any) -> List[Dict[str, Any]]: return [self._entry_to_event(entry_id, fields) for entry_id, fields in rows or []] @staticmethod def _matches_city(event: Dict[str, Any], city_set: Set[str]) -> bool: return not city_set or str(event.get("city") or "").strip().lower() in city_set def _replay_events_forward( self, *, city_set: Set[str], since: int, limit: int, ) -> tuple[List[Dict[str, Any]], bool, int]: replay: List[Dict[str, Any]] = [] scanned = 0 min_id = "-" hit_boundary = False while len(replay) < limit and scanned < self.replay_scan_max: count = min(self.replay_chunk_size, self.replay_scan_max - scanned) rows = self._client.xrange(self.stream_key, min=min_id, max="+", count=count) if not rows: hit_boundary = True break scanned += len(rows) for entry_id, fields in rows: event = self._entry_to_event(entry_id, fields) if int(event.get("revision") or 0) <= since: continue if self._matches_city(event, city_set): replay.append(event) if len(replay) >= limit: break min_id = f"({_decode(rows[-1][0])}" return replay, hit_boundary, scanned def _replay_events_reverse( self, *, city_set: Set[str], since: int, limit: int, ) -> tuple[List[Dict[str, Any]], bool, int]: replay_desc: List[Dict[str, Any]] = [] scanned = 0 max_id = "+" hit_boundary = False while len(replay_desc) < limit and scanned < self.replay_scan_max: count = min(self.replay_chunk_size, self.replay_scan_max - scanned) rows = self._client.xrevrange(self.stream_key, max=max_id, min="-", count=count) if not rows: hit_boundary = True break scanned += len(rows) for entry_id, fields in rows: event = self._entry_to_event(entry_id, fields) if int(event.get("revision") or 0) <= since: hit_boundary = True return list(reversed(replay_desc)), hit_boundary, scanned if self._matches_city(event, city_set): replay_desc.append(event) if len(replay_desc) >= limit: break max_id = f"({_decode(rows[-1][0])}" return list(reversed(replay_desc)), hit_boundary, scanned @staticmethod def _entry_to_event(entry_id: Any, fields: Dict[Any, Any]) -> Dict[str, Any]: normalized = {_decode(key): _decode(value) for key, value in dict(fields or {}).items()} payload = json.loads(normalized.get("payload_json") or "{}") schema_type = normalized.get("schema_type") or "city_observation_patch" schema_version = int(normalized.get("schema_version") or 1) created_at_ms = _int_or_zero(normalized.get("created_at_ms")) or int(time.time() * 1000) ts = _int_or_zero(normalized.get("ts")) or created_at_ms obs_time = normalized.get("obs_time") or None return { "type": normalized.get("type") or f"{schema_type}.v{schema_version}", "revision": int(normalized["revision"]), "city": normalized.get("city") or "", "source": normalized.get("source") or "", "obs_time": obs_time, **_time_contract_from_payload(payload if isinstance(payload, dict) else {}), "ts": ts, "payload": payload if isinstance(payload, dict) else {}, }