wowlab_sentinel/scheduler/
burst.rs

1//! Autoscales the Latitude.sh burst node pool based on pending-chunk queue depth.
2
3#![expect(
4    clippy::cast_possible_truncation,
5    clippy::cast_precision_loss,
6    clippy::cast_sign_loss,
7    reason = "burst sizing intentionally rounds bounded database counts through floating-point arithmetic"
8)]
9
10use std::{
11    sync::atomic::Ordering,
12    time::{Duration, Instant},
13};
14
15use async_trait::async_trait;
16use tokio::sync::Mutex;
17
18use crate::{
19    ExposeSecret,
20    cron::CronJob,
21    latitude::{LatitudeClient, LatitudeError, ProvisionSpec},
22    state::ServerState,
23};
24
25const BURST_PLAN: &str = "m4-metal-xlarge";
26const BURST_OS: &str = "ubuntu_24_04";
27const BURST_HOSTNAME: &str = "wowlab-burst";
28const SCALE_DOWN_HYSTERESIS: Duration = Duration::from_secs(15 * 60);
29const MIN_NODE_AGE_BEFORE_KILL: Duration = Duration::from_secs(55 * 60);
30
31#[derive(Clone, Debug)]
32struct BurstNode {
33    server_id: String,
34    started_at: Instant,
35    below_target_since: Option<Instant>,
36}
37
38#[derive(Debug)]
39pub(crate) struct BurstScheduler {
40    schedule: String,
41    nodes: Mutex<Vec<BurstNode>>,
42}
43
44struct ReconcileCtx<'a> {
45    latitude: &'a LatitudeClient,
46    project: &'a str,
47    site: &'a str,
48    user_data_b64: &'a str,
49    tag_ids: &'a [&'a str],
50}
51
52impl BurstScheduler {
53    pub(crate) fn new(schedule: &str) -> Self {
54        Self {
55            schedule: schedule.to_string(),
56            nodes: Mutex::new(Vec::new()),
57        }
58    }
59}
60
61fn compute_target(pending: i64, floor_capacity: i64, per_node: i64) -> usize {
62    let overflow = pending.saturating_sub(floor_capacity);
63
64    if overflow <= 0 || per_node <= 0 {
65        return 0;
66    }
67
68    ((overflow as f64) / (per_node as f64)).ceil() as usize
69}
70
71#[async_trait]
72impl CronJob for BurstScheduler {
73    fn name(&self) -> &'static str {
74        "burst_scheduler"
75    }
76
77    fn schedule(&self) -> &str {
78        &self.schedule
79    }
80
81    async fn run(&self, state: &ServerState) {
82        if !state.config.burst_enabled {
83            return;
84        }
85
86        let Some(latitude) = state.latitude.as_ref() else {
87            tracing::debug!("Burst enabled but no Latitude client configured");
88
89            return;
90        };
91        let Some(project) = state.config.latitude_project.as_deref() else {
92            tracing::warn!("Burst enabled but LATITUDE_PROJECT not set");
93
94            return;
95        };
96        let Some(site) = state.config.latitude_site.as_deref() else {
97            tracing::warn!("Burst enabled but LATITUDE_SITE not set");
98
99            return;
100        };
101        let Some(claim_token_secret) = state.config.burst_node_claim_token.as_ref() else {
102            tracing::warn!("Burst enabled but BURST_NODE_CLAIM_TOKEN not set");
103
104            return;
105        };
106        let Some(ts_authkey_secret) = state.config.burst_ts_authkey.as_ref() else {
107            tracing::warn!("Burst enabled but BURST_TS_AUTHKEY not set");
108
109            return;
110        };
111
112        let pending = state.runtime.pending_chunks.load(Ordering::Relaxed);
113        let target = compute_target(
114            pending,
115            state.config.burst_floor_capacity,
116            state.config.burst_per_node_throughput,
117        )
118        .min(state.config.burst_max_nodes);
119
120        let user_data = render_user_data(
121            BURST_HOSTNAME,
122            claim_token_secret.expose_secret(),
123            ts_authkey_secret.expose_secret(),
124        );
125        let tag_ids: Vec<&str> = state
126            .config
127            .latitude_burst_tag_id
128            .as_deref()
129            .into_iter()
130            .collect();
131
132        let ctx = ReconcileCtx {
133            latitude,
134            project,
135            site,
136            user_data_b64: &user_data,
137            tag_ids: &tag_ids,
138        };
139
140        let mut nodes = self.nodes.lock().await;
141
142        if let Err(e) = reconcile(&mut nodes, target, &ctx).await {
143            tracing::warn!(error = %e, "Burst reconcile failed");
144        }
145    }
146}
147
148async fn reconcile(
149    nodes: &mut Vec<BurstNode>,
150    target: usize,
151    ctx: &ReconcileCtx<'_>,
152) -> Result<(), LatitudeError> {
153    let now = Instant::now();
154
155    if target > nodes.len() {
156        let to_add = target - nodes.len();
157
158        for _ in 0..to_add {
159            match provision_node(ctx, now).await {
160                Ok(node) => nodes.push(node),
161                Err(error) => {
162                    if let Some(server_id) = error.stranded_server_id() {
163                        nodes.push(BurstNode {
164                            server_id: server_id.to_owned(),
165                            started_at: now,
166                            below_target_since: None,
167                        });
168                    }
169
170                    return Err(error);
171                }
172            }
173        }
174
175        return Ok(());
176    }
177
178    if target < nodes.len() {
179        mark_below_target(nodes, now);
180        let needed = nodes.len() - target;
181        let killable: Vec<String> = nodes
182            .iter()
183            .filter(|n| {
184                n.below_target_since
185                    .is_some_and(|t| now.duration_since(t) > SCALE_DOWN_HYSTERESIS)
186                    && now.duration_since(n.started_at) > MIN_NODE_AGE_BEFORE_KILL
187            })
188            .take(needed)
189            .map(|n| n.server_id.clone())
190            .collect();
191
192        let mut deprovisioned = Vec::with_capacity(killable.len());
193
194        for id in killable {
195            if deprovision_node(ctx.latitude, &id).await.is_ok() {
196                deprovisioned.push(id);
197            }
198        }
199
200        nodes.retain(|n| !deprovisioned.contains(&n.server_id));
201
202        return Ok(());
203    }
204
205    clear_below_target(nodes);
206
207    Ok(())
208}
209
210async fn provision_node(ctx: &ReconcileCtx<'_>, now: Instant) -> Result<BurstNode, LatitudeError> {
211    let spec = ProvisionSpec {
212        project: ctx.project,
213        plan: BURST_PLAN,
214        site: ctx.site,
215        operating_system: BURST_OS,
216        hostname: BURST_HOSTNAME,
217        user_data_b64: ctx.user_data_b64,
218        tag_ids: ctx.tag_ids,
219    };
220    let server = ctx.latitude.provision(&spec).await?;
221
222    tracing::info!(server_id = %server.id, "Provisioned burst node");
223
224    Ok(BurstNode {
225        server_id: server.id,
226        started_at: now,
227        below_target_since: None,
228    })
229}
230
231fn mark_below_target(nodes: &mut [BurstNode], now: Instant) {
232    for node in nodes {
233        if node.below_target_since.is_none() {
234            node.below_target_since = Some(now);
235        }
236    }
237}
238
239async fn deprovision_node(latitude: &LatitudeClient, server_id: &str) -> Result<(), LatitudeError> {
240    match latitude.deprovision(server_id).await {
241        Ok(()) => {
242            tracing::info!(server_id, "Deprovisioned burst node");
243
244            Ok(())
245        }
246        Err(error) => {
247            tracing::warn!(server_id, %error, "Burst node deprovision failed");
248
249            Err(error)
250        }
251    }
252}
253
254fn clear_below_target(nodes: &mut [BurstNode]) {
255    for node in nodes {
256        node.below_target_since = None;
257    }
258}
259
260// docref:start deployment-render-user-data
261fn render_user_data(hostname: &str, claim_token: &str, ts_authkey: &str) -> String {
262    use base64::Engine;
263    let template = include_str!("../../../../deploy/node/latitude-userdata.yaml.tmpl");
264    let rendered = template
265        .replace("__NODE_NAME__", hostname)
266        .replace("__CLAIM_TOKEN__", claim_token)
267        .replace("__TS_AUTHKEY__", ts_authkey);
268    base64::engine::general_purpose::STANDARD.encode(rendered.as_bytes())
269}
270// docref:end deployment-render-user-data
271
272#[cfg(test)]
273mod tests {
274    use base64::Engine;
275    use googletest::prelude::*;
276    use secrecy::SecretString;
277    use wiremock::{
278        Mock, MockServer, ResponseTemplate,
279        matchers::{method, path},
280    };
281
282    use super::*;
283
284    #[gtest]
285    fn target_zero_when_pending_below_floor() -> Result<()> {
286        verify_eq!(compute_target(100, 200, 50), 0)?;
287
288        Ok(())
289    }
290
291    #[gtest]
292    fn target_ceil_of_overflow_over_per_node() -> Result<()> {
293        verify_eq!(compute_target(250, 100, 50), 3)?;
294        verify_eq!(compute_target(251, 100, 50), 4)?;
295
296        Ok(())
297    }
298
299    #[gtest]
300    fn target_zero_when_per_node_zero() -> Result<()> {
301        verify_eq!(compute_target(500, 0, 0), 0)?;
302
303        Ok(())
304    }
305
306    #[gtest]
307    fn user_data_resolves_every_cloud_init_placeholder() -> Result<()> {
308        let encoded = render_user_data("burst-01", "claim-secret", "tailnet-secret");
309        let bytes = base64::engine::general_purpose::STANDARD
310            .decode(encoded)
311            .or_fail()?;
312        let rendered = String::from_utf8(bytes).or_fail()?;
313
314        verify_true!(rendered.contains("NODE_NAME=burst-01"))?;
315        verify_true!(rendered.contains("NODE_CLAIM_TOKEN=claim-secret"))?;
316        verify_true!(rendered.contains("--authkey=tailnet-secret"))?;
317        verify_true!(!rendered.contains("__NODE_NAME__"))?;
318        verify_true!(!rendered.contains("__CLAIM_TOKEN__"))?;
319        verify_true!(!rendered.contains("__TS_AUTHKEY__"))?;
320
321        Ok(())
322    }
323
324    #[gtest]
325    #[tokio::test]
326    async fn failed_deprovision_keeps_node_tracked_for_retry() -> Result<()> {
327        let mock = MockServer::start().await;
328
329        Mock::given(method("DELETE"))
330            .and(path("/servers/sv_retry"))
331            .respond_with(ResponseTemplate::new(500))
332            .expect(1)
333            .mount(&mock)
334            .await;
335
336        let latitude = LatitudeClient::new(SecretString::from("test-token"), mock.uri());
337        let now = Instant::now();
338        let old_enough = MIN_NODE_AGE_BEFORE_KILL + Duration::from_secs(1);
339        let below_long_enough = SCALE_DOWN_HYSTERESIS + Duration::from_secs(1);
340        let started_at = now.checked_sub(old_enough).or_fail()?;
341        let below_target_since = now.checked_sub(below_long_enough).or_fail()?;
342        let mut nodes = vec![BurstNode {
343            server_id: "sv_retry".to_string(),
344            started_at,
345            below_target_since: Some(below_target_since),
346        }];
347        let ctx = ReconcileCtx {
348            latitude: &latitude,
349            project: "project",
350            site: "site",
351            user_data_b64: "user-data",
352            tag_ids: &[],
353        };
354
355        reconcile(&mut nodes, 0, &ctx).await.or_fail()?;
356
357        verify_that!(nodes, len(eq(1)))?;
358
359        verify_that!(nodes[0].server_id.as_str(), eq("sv_retry"))
360    }
361
362    #[gtest]
363    #[tokio::test]
364    async fn failed_provision_cleanup_keeps_created_server_tracked() -> Result<()> {
365        let mock = MockServer::start().await;
366
367        Mock::given(method("POST"))
368            .and(path("/servers"))
369            .respond_with(ResponseTemplate::new(201).set_body_json(serde_json::json!({
370                "data": { "type": "servers", "id": "sv_stranded", "attributes": {} }
371            })))
372            .mount(&mock)
373            .await;
374        Mock::given(method("PATCH"))
375            .and(path("/servers/sv_stranded"))
376            .respond_with(ResponseTemplate::new(500))
377            .mount(&mock)
378            .await;
379        Mock::given(method("DELETE"))
380            .and(path("/servers/sv_stranded"))
381            .respond_with(ResponseTemplate::new(500))
382            .mount(&mock)
383            .await;
384
385        let latitude = LatitudeClient::new(SecretString::from("test-token"), mock.uri());
386        let mut nodes = Vec::new();
387        let ctx = ReconcileCtx {
388            latitude: &latitude,
389            project: "project",
390            site: "site",
391            user_data_b64: "user-data",
392            tag_ids: &["tag"],
393        };
394
395        verify_that!(reconcile(&mut nodes, 1, &ctx).await, err(anything()))?;
396        verify_that!(nodes, len(eq(1)))?;
397
398        verify_that!(nodes[0].server_id.as_str(), eq("sv_stranded"))
399    }
400}
wowlab_sentinel/scheduler/burst.rs

wowlab_sentinel/scheduler/
burst.rs
