1use std::{
2 sync::{
3 Arc,
4 atomic::{AtomicI64, AtomicU64, Ordering},
5 },
6 time::Duration,
7};
8
9use metrics_exporter_prometheus::PrometheusHandle;
10use serde::Serialize;
11use wowlab_centrifuge::{Client as CentrifugeClient, Presence};
12use wowlab_common::time::{self, Instant};
13use wowlab_types::sim::FastMap;
14
15use crate::{config::Config, utils::filter_refresh::FilterMap};
16
17#[derive(Debug, Default)]
19pub(crate) struct McpCache(FastMap<&'static str, String>);
20
21impl McpCache {
22 pub(crate) fn new() -> Self {
23 Self::default()
24 }
25
26 pub(crate) fn get(&self, key: &str) -> Option<&str> {
27 self.0.get(key).map(String::as_str)
28 }
29
30 pub(crate) fn set(&mut self, key: &'static str, value: String) {
31 self.0.insert(key, value);
32 }
33}
34
35#[derive(Debug)]
36pub(crate) struct McpState {
37 docs: tokio::sync::RwLock<McpCache>,
38 pub schema: crate::mcp::schema::SchemaCatalog,
39}
40
41impl McpState {
42 pub(crate) fn new() -> Result<Self, crate::mcp::schema::SchemaCatalogError> {
43 Ok(Self {
44 docs: tokio::sync::RwLock::new(McpCache::new()),
45 schema: crate::mcp::schema::SchemaCatalog::try_new()?,
46 })
47 }
48
49 pub(crate) async fn docs(&self) -> tokio::sync::RwLockReadGuard<'_, McpCache> {
50 self.docs.read().await
51 }
52
53 pub(crate) async fn docs_mut(&self) -> tokio::sync::RwLockWriteGuard<'_, McpCache> {
54 self.docs.write().await
55 }
56}
57
58pub(crate) type McpStateHandle = Arc<McpState>;
59
60const PUBLISH_RETRIES: u32 = 3;
61const PUBLISH_RETRY_DELAY: Duration = Duration::from_millis(100);
62const SCHEDULER_GRACE_PERIOD_SECS: u64 = 60;
63const PRESENCE_GRACE_PERIOD_SECS: u64 = 90;
64const CRON_GRACE_PERIOD_SECS: u64 = 90;
65
66#[derive(Debug)]
67pub(crate) struct HealthState {
68 pub started_at: Instant,
69 pub prometheus: PrometheusHandle,
70 last_scheduler_tick: AtomicU64,
71 last_presence_tick: AtomicU64,
72 last_cron_tick: AtomicU64,
73}
74
75impl HealthState {
76 pub(crate) fn new(prometheus: PrometheusHandle) -> Self {
77 Self {
78 started_at: Instant::now(),
79 prometheus,
80 last_scheduler_tick: AtomicU64::new(0),
81 last_presence_tick: AtomicU64::new(0),
82 last_cron_tick: AtomicU64::new(0),
83 }
84 }
85}
86
87#[derive(Debug, Default)]
88pub(crate) struct RuntimeState {
89 pub jobs: tokio::sync::RwLock<crate::scheduler::runtime::JobRuntimeStore>,
90 pub pending_chunks: AtomicI64,
91}
92
93pub(crate) struct ServerState {
94 pub config: Config,
95 pub dbs: crate::db::DbRegistry,
96 pub filters: FilterMap,
97 pub health: HealthState,
98 pub runtime: RuntimeState,
99 pub centrifuge: CentrifugeClient,
100 pub presence: Presence,
101 pub mcp: McpStateHandle,
102 pub latitude: Option<crate::latitude::LatitudeClient>,
103}
104
105impl ServerState {
106 pub(crate) fn touch_scheduler(&self) {
107 touch(&self.health.last_scheduler_tick);
108 }
109
110 pub(crate) fn touch_presence(&self) {
111 touch(&self.health.last_presence_tick);
112 }
113
114 pub(crate) fn touch_cron(&self) {
115 touch(&self.health.last_cron_tick);
116 }
117
118 pub(crate) fn scheduler_healthy(&self) -> bool {
119 is_healthy(
120 &self.health.last_scheduler_tick,
121 SCHEDULER_GRACE_PERIOD_SECS,
122 )
123 }
124
125 pub(crate) fn presence_healthy(&self) -> bool {
126 is_healthy(&self.health.last_presence_tick, PRESENCE_GRACE_PERIOD_SECS)
127 }
128
129 pub(crate) fn cron_healthy(&self) -> bool {
130 is_healthy(&self.health.last_cron_tick, CRON_GRACE_PERIOD_SECS)
131 }
132
133 pub(crate) async fn centrifuge_healthy(&self) -> bool {
134 self.centrifuge.is_connected().await
135 }
136
137 pub(crate) async fn publish<T>(&self, channel: &str, payload: &T)
138 where
139 T: Serialize,
140 {
141 let data = match serde_json::to_vec(payload) {
142 Ok(d) => d,
143 Err(e) => {
144 tracing::error!(error = %e, channel, "Failed to serialize publish payload");
145
146 return;
147 }
148 };
149
150 let mut last_error = None;
152 for attempt in 0..PUBLISH_RETRIES {
153 match self.centrifuge.publish(channel, data.clone()).await {
155 Ok(()) => return,
156 Err(error) if error.is_temporary() && attempt + 1 < PUBLISH_RETRIES => {
157 last_error = Some(error);
158 tokio::time::sleep(PUBLISH_RETRY_DELAY).await;
159 }
160 Err(error) => {
161 last_error = Some(error);
162 break;
163 }
164 }
165 }
166 if let Some(error) = last_error {
167 tracing::warn!(%error, channel, "Failed to publish");
168 }
169 }
171}
172
173impl std::fmt::Debug for ServerState {
174 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
175 f.debug_struct("ServerState")
176 .field("config", &self.config)
177 .field("dbs", &self.dbs)
178 .field("filters", &"<FilterMap>")
179 .field("health", &self.health)
180 .field("runtime", &self.runtime)
181 .field("centrifuge", &"<CentrifugeClient>")
182 .field("presence", &"<Presence>")
183 .field("mcp", &self.mcp)
184 .field(
185 "latitude",
186 &self.latitude.as_ref().map(|_| "<LatitudeClient>"),
187 )
188 .finish()
189 }
190}
191
192fn epoch_secs() -> u64 {
193 time::unix_timestamp_secs()
194}
195
196fn touch(tick: &AtomicU64) {
197 tick.store(epoch_secs(), Ordering::Relaxed);
198}
199
200fn is_healthy(tick: &AtomicU64, grace_period_secs: u64) -> bool {
201 let ts = tick.load(Ordering::Relaxed);
202
203 epoch_secs().saturating_sub(ts) < grace_period_secs
204}