bhtune_server/
state.rs

1//! [`AppState`]: the shared state every route handler receives via `axum::extract::State`.
2
3use bhtune_db::SqlitePool;
4use chrono::{DateTime, Duration, Utc};
5use std::collections::{HashMap, VecDeque};
6use std::sync::{Arc, RwLock};
7
8use crate::active_run::ActiveRun;
9use bhtune_runtime::config::{DemoPolicy, ServerMode};
10use tokio::sync::Mutex;
11use tokio::sync::{OwnedMutexGuard, OwnedSemaphorePermit, Semaphore};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub struct DemoQuotaExceeded {
15    pub retry_after_secs: u64,
16}
17
18#[derive(Debug, Clone)]
19struct AcceptedStart {
20    id: u64,
21    accepted_at: DateTime<Utc>,
22}
23
24#[derive(Default)]
25struct AcceptedStartWindows {
26    next_id: u64,
27    by_token: HashMap<String, VecDeque<AcceptedStart>>,
28    by_ip: HashMap<String, VecDeque<AcceptedStart>>,
29}
30
31#[derive(Debug)]
32pub struct DemoAcceptedStart {
33    id: u64,
34    token_hash: String,
35    client_ip: String,
36}
37
38#[derive(Clone)]
39pub struct DemoRuntime {
40    run_permits: Arc<Semaphore>,
41    sse_permits: Arc<Semaphore>,
42    ordinary_request_permits: Arc<Semaphore>,
43    start_admission: Arc<Mutex<()>>,
44    accepted_starts: Arc<Mutex<AcceptedStartWindows>>,
45    visitor_run_permits: Arc<Mutex<HashMap<String, Arc<Semaphore>>>>,
46    visitor_sse_permits: Arc<Mutex<HashMap<String, Arc<Semaphore>>>>,
47}
48
49impl DemoRuntime {
50    fn new(policy: DemoPolicy) -> Self {
51        Self {
52            run_permits: Arc::new(Semaphore::new(policy.max_active_runs_global as usize)),
53            sse_permits: Arc::new(Semaphore::new(policy.max_sse_global as usize)),
54            ordinary_request_permits: Arc::new(Semaphore::new(
55                policy.ordinary_request_concurrency as usize,
56            )),
57            start_admission: Arc::new(Mutex::new(())),
58            accepted_starts: Arc::new(Mutex::new(AcceptedStartWindows::default())),
59            visitor_run_permits: Arc::new(Mutex::new(HashMap::new())),
60            visitor_sse_permits: Arc::new(Mutex::new(HashMap::new())),
61        }
62    }
63
64    pub fn try_acquire_global_run(
65        &self,
66    ) -> Result<OwnedSemaphorePermit, tokio::sync::TryAcquireError> {
67        self.run_permits.clone().try_acquire_owned()
68    }
69
70    pub fn try_acquire_global_sse(
71        &self,
72    ) -> Result<OwnedSemaphorePermit, tokio::sync::TryAcquireError> {
73        self.sse_permits.clone().try_acquire_owned()
74    }
75
76    pub fn try_acquire_ordinary_request(
77        &self,
78    ) -> Result<OwnedSemaphorePermit, tokio::sync::TryAcquireError> {
79        self.ordinary_request_permits.clone().try_acquire_owned()
80    }
81
82    /// Serializes the short admission phase that checks global row capacity and persists a
83    /// prepared run. Without this guard, several simultaneous starts could all observe one
84    /// remaining row slot before any of them inserted, exceeding the fixed global Demo cap.
85    pub async fn lock_start_admission(&self) -> OwnedMutexGuard<()> {
86        self.start_admission.clone().lock_owned().await
87    }
88
89    pub async fn try_acquire_visitor_run(
90        &self,
91        token_hash: &str,
92        limit: u32,
93    ) -> Result<OwnedSemaphorePermit, tokio::sync::TryAcquireError> {
94        visitor_semaphore(&self.visitor_run_permits, token_hash, limit)
95            .await
96            .try_acquire_owned()
97    }
98
99    pub async fn try_acquire_visitor_sse(
100        &self,
101        token_hash: &str,
102        limit: u32,
103    ) -> Result<OwnedSemaphorePermit, tokio::sync::TryAcquireError> {
104        visitor_semaphore(&self.visitor_sse_permits, token_hash, limit)
105            .await
106            .try_acquire_owned()
107    }
108
109    /// Atomically reserves one accepted-start slot in both the visitor-token and client-IP
110    /// windows. A caller that fails before scheduling the run must release the returned handle;
111    /// a successfully scheduled run leaves it recorded until the fixed window expires.
112    pub async fn reserve_accepted_start(
113        &self,
114        token_hash: &str,
115        client_ip: &str,
116        now: DateTime<Utc>,
117        policy: DemoPolicy,
118    ) -> Result<DemoAcceptedStart, DemoQuotaExceeded> {
119        let mut windows = self.accepted_starts.lock().await;
120        prune_start_windows(&mut windows, now, policy.accepted_start_window_secs);
121        if let Some(retry_after_secs) =
122            quota_retry_after(&windows, token_hash, client_ip, now, policy)
123        {
124            return Err(DemoQuotaExceeded { retry_after_secs });
125        }
126
127        windows.next_id = windows.next_id.wrapping_add(1);
128        let accepted = AcceptedStart {
129            id: windows.next_id,
130            accepted_at: now,
131        };
132        windows
133            .by_token
134            .entry(token_hash.to_owned())
135            .or_default()
136            .push_back(accepted.clone());
137        windows
138            .by_ip
139            .entry(client_ip.to_owned())
140            .or_default()
141            .push_back(accepted.clone());
142        Ok(DemoAcceptedStart {
143            id: accepted.id,
144            token_hash: token_hash.to_owned(),
145            client_ip: client_ip.to_owned(),
146        })
147    }
148
149    pub async fn release_accepted_start(&self, accepted: DemoAcceptedStart) {
150        let mut windows = self.accepted_starts.lock().await;
151        remove_accepted_start(&mut windows.by_token, &accepted.token_hash, accepted.id);
152        remove_accepted_start(&mut windows.by_ip, &accepted.client_ip, accepted.id);
153    }
154
155    /// Reclaims expired quota entries and inactive per-visitor semaphores during the periodic
156    /// Demo cleanup pass, including tokens that never make another request.
157    pub async fn cleanup(&self, now: DateTime<Utc>, policy: DemoPolicy) {
158        {
159            let mut accepted_starts = self.accepted_starts.lock().await;
160            prune_start_windows(&mut accepted_starts, now, policy.accepted_start_window_secs);
161        }
162        {
163            let mut visitor_runs = self.visitor_run_permits.lock().await;
164            retain_active_visitor_semaphores(&mut visitor_runs, policy.max_active_runs_per_visitor);
165        }
166        {
167            let mut visitor_streams = self.visitor_sse_permits.lock().await;
168            retain_active_visitor_semaphores(&mut visitor_streams, policy.max_sse_per_visitor);
169        }
170    }
171}
172
173async fn visitor_semaphore(
174    permits: &Mutex<HashMap<String, Arc<Semaphore>>>,
175    token_hash: &str,
176    limit: u32,
177) -> Arc<Semaphore> {
178    let mut permits = permits.lock().await;
179    permits.retain(|_, permit| {
180        Arc::strong_count(permit) > 1 || permit.available_permits() < limit as usize
181    });
182    permits
183        .entry(token_hash.to_owned())
184        .or_insert_with(|| Arc::new(Semaphore::new(limit as usize)))
185        .clone()
186}
187
188fn retain_active_visitor_semaphores(permits: &mut HashMap<String, Arc<Semaphore>>, limit: u32) {
189    permits.retain(|_, permit| {
190        Arc::strong_count(permit) > 1 || permit.available_permits() < limit as usize
191    });
192}
193
194fn prune_start_windows(windows: &mut AcceptedStartWindows, now: DateTime<Utc>, window_secs: u64) {
195    let cutoff = now - Duration::seconds(window_secs as i64);
196    windows.by_token.retain(|_, starts| {
197        while starts
198            .front()
199            .is_some_and(|start| start.accepted_at <= cutoff)
200        {
201            starts.pop_front();
202        }
203        !starts.is_empty()
204    });
205    windows.by_ip.retain(|_, starts| {
206        while starts
207            .front()
208            .is_some_and(|start| start.accepted_at <= cutoff)
209        {
210            starts.pop_front();
211        }
212        !starts.is_empty()
213    });
214}
215
216fn quota_retry_after(
217    windows: &AcceptedStartWindows,
218    token_hash: &str,
219    client_ip: &str,
220    now: DateTime<Utc>,
221    policy: DemoPolicy,
222) -> Option<u64> {
223    let token_retry = retry_after_for(
224        windows.by_token.get(token_hash),
225        policy.accepted_starts_per_token,
226        now,
227        policy.accepted_start_window_secs,
228    );
229    let ip_retry = retry_after_for(
230        windows.by_ip.get(client_ip),
231        policy.accepted_starts_per_client_ip,
232        now,
233        policy.accepted_start_window_secs,
234    );
235    match (token_retry, ip_retry) {
236        (Some(left), Some(right)) => Some(left.max(right)),
237        (Some(retry), None) | (None, Some(retry)) => Some(retry),
238        (None, None) => None,
239    }
240}
241
242fn retry_after_for(
243    starts: Option<&VecDeque<AcceptedStart>>,
244    limit: u32,
245    now: DateTime<Utc>,
246    window_secs: u64,
247) -> Option<u64> {
248    let starts = starts?;
249    if starts.len() < limit as usize {
250        return None;
251    }
252    let remaining_ms = (starts.front()?.accepted_at + Duration::seconds(window_secs as i64) - now)
253        .num_milliseconds()
254        .max(1);
255    Some((remaining_ms as u64).div_ceil(1_000))
256}
257
258fn remove_accepted_start(
259    windows: &mut HashMap<String, VecDeque<AcceptedStart>>,
260    key: &str,
261    id: u64,
262) {
263    let Some(starts) = windows.get_mut(key) else {
264        return;
265    };
266    starts.retain(|start| start.id != id);
267    if starts.is_empty() {
268        windows.remove(key);
269    }
270}
271
272/// Cheap to clone (an `Arc`-backed connection pool under the hood, per `sqlx`, and
273/// [`ActiveRun`] is itself `Arc<Mutex<..>>`-backed) -- axum's `State` extractor just requires
274/// `Clone` to hand a copy to each request.
275#[derive(Clone)]
276pub struct AppState {
277    pub pool: SqlitePool,
278    /// Read-only connection used by the Full-mode preflight endpoint to resolve templates.
279    pub(crate) preflight_database: Option<SqlitePool>,
280    /// The in-flight tune registry and exclusive post-hoc write/revert reservation -- see
281    /// [`ActiveRun`]'s own doc comment for the concurrency and shutdown behavior.
282    pub active_run: ActiveRun,
283    /// The live, revisioned TOML configuration. Route handlers take a fresh snapshot for
284    /// every operation so a configuration-page save is visible without restarting the server.
285    pub config_store: Arc<RwLock<bhtune_runtime::config::LoadedConfigStore>>,
286    pub allowed_origin: Option<String>,
287    pub trusted_proxy: Option<String>,
288    pub mode: ServerMode,
289    pub demo_policy: DemoPolicy,
290    pub demo_runtime: DemoRuntime,
291}
292
293impl AppState {
294    pub fn config_snapshot(&self) -> anyhow::Result<bhtune_runtime::config::BhtuneConfig> {
295        self.config_store
296            .read()
297            .map(|store| store.config.clone())
298            .map_err(|_| anyhow::anyhow!("configuration store lock is poisoned"))
299    }
300
301    pub fn for_mode(
302        pool: SqlitePool,
303        config_store: Arc<RwLock<bhtune_runtime::config::LoadedConfigStore>>,
304        mode: ServerMode,
305        demo_policy: DemoPolicy,
306    ) -> Self {
307        Self::for_mode_with_network_config(pool, config_store, mode, demo_policy, None, None)
308    }
309
310    pub fn for_mode_with_network_config(
311        pool: SqlitePool,
312        config_store: Arc<RwLock<bhtune_runtime::config::LoadedConfigStore>>,
313        mode: ServerMode,
314        demo_policy: DemoPolicy,
315        allowed_origin: Option<String>,
316        trusted_proxy: Option<String>,
317    ) -> Self {
318        Self {
319            pool,
320            preflight_database: None,
321            active_run: ActiveRun::default(),
322            config_store,
323            allowed_origin,
324            trusted_proxy,
325            demo_runtime: DemoRuntime::new(demo_policy),
326            mode,
327            demo_policy,
328        }
329    }
330
331    pub(crate) fn with_preflight_database(
332        mut self,
333        preflight_database: Option<SqlitePool>,
334    ) -> Self {
335        self.preflight_database = preflight_database;
336        self
337    }
338}
339
340#[cfg(test)]
341mod tests {
342    use super::*;
343
344    #[tokio::test]
345    async fn config_snapshot_returns_an_independent_config_copy() {
346        let state = crate::test_support::in_memory_state().await;
347        let mut snapshot = state.config_snapshot().unwrap();
348        snapshot.allow_uncertain_quality = false;
349
350        assert!(state.config_snapshot().unwrap().allow_uncertain_quality);
351    }
352
353    #[tokio::test]
354    async fn config_snapshot_reports_a_poisoned_store_lock() {
355        let state = crate::test_support::in_memory_state().await;
356        let store = state.config_store.clone();
357        let _ = std::thread::spawn(move || {
358            let _guard = store.write().unwrap();
359            panic!("deliberately poison the test lock");
360        })
361        .join();
362
363        assert!(
364            state
365                .config_snapshot()
366                .unwrap_err()
367                .to_string()
368                .contains("poisoned")
369        );
370    }
371
372    #[tokio::test]
373    async fn demo_runtime_uses_the_global_and_ordinary_policy_limits() {
374        let policy = DemoPolicy {
375            max_active_runs_global: 1,
376            max_sse_global: 1,
377            ordinary_request_concurrency: 1,
378            ..DemoPolicy::default()
379        };
380        let runtime = DemoRuntime::new(policy);
381
382        let _run = runtime.try_acquire_global_run().unwrap();
383        assert!(runtime.try_acquire_global_run().is_err());
384        let _sse = runtime.try_acquire_global_sse().unwrap();
385        assert!(runtime.try_acquire_global_sse().is_err());
386        let _request = runtime.try_acquire_ordinary_request().unwrap();
387        assert!(runtime.try_acquire_ordinary_request().is_err());
388    }
389
390    #[tokio::test]
391    async fn demo_start_admission_is_serialized() {
392        let runtime = DemoRuntime::new(DemoPolicy::default());
393        let admission = runtime.lock_start_admission().await;
394        let waiting_runtime = runtime.clone();
395        let waiter = tokio::spawn(async move {
396            let _admission = waiting_runtime.lock_start_admission().await;
397        });
398        tokio::task::yield_now().await;
399        assert!(!waiter.is_finished());
400
401        drop(admission);
402
403        waiter.await.unwrap();
404    }
405
406    #[tokio::test]
407    async fn accepted_start_quotas_are_independent_for_token_and_client_ip() {
408        let policy = DemoPolicy::default();
409        let now = DateTime::from_timestamp(1_700_000_000, 0).unwrap();
410        let token_runtime = DemoRuntime::new(policy);
411        for index in 0..policy.accepted_starts_per_token {
412            token_runtime
413                .reserve_accepted_start("token", &format!("ip-{index}"), now, policy)
414                .await
415                .unwrap();
416        }
417        let token_error = token_runtime
418            .reserve_accepted_start("token", "unused-ip", now, policy)
419            .await
420            .unwrap_err();
421        assert_eq!(
422            token_error.retry_after_secs,
423            policy.accepted_start_window_secs
424        );
425        assert!(
426            token_runtime
427                .reserve_accepted_start(
428                    "token",
429                    "unused-ip",
430                    now + Duration::seconds(policy.accepted_start_window_secs as i64),
431                    policy,
432                )
433                .await
434                .is_ok()
435        );
436
437        let ip_runtime = DemoRuntime::new(policy);
438        for index in 0..policy.accepted_starts_per_client_ip {
439            ip_runtime
440                .reserve_accepted_start(&format!("token-{index}"), "shared-ip", now, policy)
441                .await
442                .unwrap();
443        }
444        assert!(
445            ip_runtime
446                .reserve_accepted_start("unused-token", "shared-ip", now, policy)
447                .await
448                .is_err()
449        );
450        assert!(
451            ip_runtime
452                .reserve_accepted_start("unused-token", "unused-ip", now, policy)
453                .await
454                .is_ok()
455        );
456    }
457
458    #[tokio::test]
459    async fn accepted_start_quota_uses_the_larger_token_or_ip_retry_after() {
460        let policy = DemoPolicy {
461            accepted_starts_per_token: 1,
462            accepted_starts_per_client_ip: 1,
463            accepted_start_window_secs: 20,
464            ..DemoPolicy::default()
465        };
466        let now = DateTime::from_timestamp(1_700_000_000, 0).unwrap();
467        let runtime = DemoRuntime::new(policy);
468
469        runtime
470            .reserve_accepted_start("token", "other-ip", now - Duration::seconds(5), policy)
471            .await
472            .unwrap();
473        runtime
474            .reserve_accepted_start(
475                "other-token",
476                "shared-ip",
477                now - Duration::seconds(10),
478                policy,
479            )
480            .await
481            .unwrap();
482
483        let error = runtime
484            .reserve_accepted_start("token", "shared-ip", now, policy)
485            .await
486            .unwrap_err();
487
488        assert_eq!(error.retry_after_secs, 15);
489    }
490
491    #[tokio::test]
492    async fn rejected_start_reservations_can_be_released_without_consuming_quota() {
493        let policy = DemoPolicy::default();
494        let now = DateTime::from_timestamp(1_700_000_000, 0).unwrap();
495        let runtime = DemoRuntime::new(policy);
496        let mut starts = Vec::new();
497        for index in 0..policy.accepted_starts_per_token {
498            starts.push(
499                runtime
500                    .reserve_accepted_start("token", &format!("ip-{index}"), now, policy)
501                    .await
502                    .unwrap(),
503            );
504        }
505        assert!(
506            runtime
507                .reserve_accepted_start("token", "another-ip", now, policy)
508                .await
509                .is_err()
510        );
511
512        runtime.release_accepted_start(starts.remove(0)).await;
513
514        assert!(
515            runtime
516                .reserve_accepted_start("token", "another-ip", now, policy)
517                .await
518                .is_ok()
519        );
520    }
521
522    #[tokio::test]
523    async fn releasing_an_accepted_start_removes_empty_token_and_ip_windows() {
524        let policy = DemoPolicy::default();
525        let now = DateTime::from_timestamp(1_700_000_000, 0).unwrap();
526        let runtime = DemoRuntime::new(policy);
527        let accepted = runtime
528            .reserve_accepted_start("token", "ip", now, policy)
529            .await
530            .unwrap();
531
532        runtime.release_accepted_start(accepted).await;
533
534        let windows = runtime.accepted_starts.lock().await;
535        assert!(windows.by_token.is_empty());
536        assert!(windows.by_ip.is_empty());
537    }
538
539    #[test]
540    fn removing_an_unknown_accepted_start_is_a_no_op() {
541        let mut windows = HashMap::new();
542
543        remove_accepted_start(&mut windows, "missing", 1);
544
545        assert!(windows.is_empty());
546    }
547
548    #[tokio::test]
549    async fn cleanup_reclaims_expired_windows_and_idle_visitor_semaphores() {
550        let policy = DemoPolicy {
551            accepted_starts_per_token: 1,
552            accepted_starts_per_client_ip: 1,
553            accepted_start_window_secs: 10,
554            max_active_runs_per_visitor: 1,
555            max_sse_per_visitor: 1,
556            ..DemoPolicy::default()
557        };
558        let runtime = DemoRuntime::new(policy);
559        let now = DateTime::from_timestamp(1_700_000_000, 0).unwrap();
560        runtime
561            .reserve_accepted_start("token", "ip", now, policy)
562            .await
563            .unwrap();
564        drop(runtime.try_acquire_visitor_run("token", 1).await.unwrap());
565        drop(runtime.try_acquire_visitor_sse("token", 1).await.unwrap());
566
567        runtime.cleanup(now + Duration::seconds(10), policy).await;
568
569        assert!(
570            runtime
571                .reserve_accepted_start("token", "ip", now + Duration::seconds(10), policy)
572                .await
573                .is_ok()
574        );
575        assert!(runtime.visitor_run_permits.lock().await.is_empty());
576        assert!(runtime.visitor_sse_permits.lock().await.is_empty());
577    }
578}
bhtune_server/state.rs

bhtune_server/
state.rs
