cronoq_server/
state.rs

1//! Shared server state.
2
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::sync::{Arc, OnceLock};
5use std::time::Duration;
6
7use bytes::Bytes;
8use rustc_hash::FxHashMap;
9use tokio_util::sync::CancellationToken;
10
11use crate::auth::Auth;
12use crate::config::LimitsConfig;
13use crate::info::TaskInfoCache;
14use crate::remote::{OutputCache, WorkerRegistry};
15
16/// State that exists for the whole process lifetime, including while standing by.
17pub struct Shared {
18    /// Credentials.
19    pub auth: Auth,
20    /// Request limits.
21    pub limits: LimitsConfig,
22    /// Running without auth on loopback: HTTP requests must name a loopback host, so a web
23    /// page cannot reach the API by rebinding its own name to 127.0.0.1.
24    pub loopback_hosts_only: bool,
25    /// True while this instance is the leader and serving.
26    pub ready: Arc<AtomicBool>,
27    /// Set once leadership is acquired and the scheduler is running.
28    pub app: OnceLock<Arc<AppState>>,
29    /// Cancelled when the server stops serving (ends Watch streams, listeners).
30    pub stop: CancellationToken,
31}
32
33impl Shared {
34    /// The serving state, or `None` while standing by (or after stopping).
35    pub fn app(&self) -> Option<&Arc<AppState>> {
36        if self.stop.is_cancelled() || !self.ready.load(Ordering::SeqCst) {
37            return None;
38        }
39        self.app.get()
40    }
41
42    /// Running and not leading: waiting for leadership, or deposed. False while shutting down.
43    pub fn is_standby(&self) -> bool {
44        !self.stop.is_cancelled() && !self.ready.load(Ordering::SeqCst)
45    }
46
47    /// The serving state as a gRPC result.
48    pub fn app_or_unavailable(&self) -> Result<&Arc<AppState>, tonic::Status> {
49        self.app()
50            .ok_or_else(|| tonic::Status::unavailable("not the leader (standby or shutting down)"))
51    }
52}
53
54/// Per-pool server settings.
55#[derive(Clone, Debug)]
56pub struct PoolInfo {
57    /// Slots a single worker may register.
58    pub slot_cap: u32,
59    /// Lease TTL workers must heartbeat within.
60    pub lease_ttl: Duration,
61    /// Attempts allowed by the pool's default retry policy.
62    pub max_attempts: u32,
63}
64
65/// State available while serving as the leader.
66pub struct AppState {
67    /// The scheduler.
68    pub sched: cronoq::Scheduler<Bytes, Bytes>,
69    /// Connected workers.
70    pub registry: Arc<WorkerRegistry>,
71    /// Recent outputs.
72    pub outputs: Arc<OutputCache>,
73    /// Recent task details.
74    pub info: Arc<TaskInfoCache>,
75    /// Pools by name.
76    pub pools: FxHashMap<Arc<str>, PoolInfo>,
77    /// Request limits.
78    pub limits: LimitsConfig,
79    /// Durable store, for outputs and details no longer cached.
80    pub queries: Option<Arc<dyn cronoq::StoreQueries<Bytes>>>,
81    /// The shard layout, when pools are split between shards.
82    pub shards: Option<Arc<crate::shard::ShardMap>>,
83    /// Connections to the other shards.
84    pub peers: Option<Arc<crate::peer::Peers>>,
85}
86
87impl AppState {
88    /// Output of a finished task: from the cache, else from the database.
89    pub async fn output(&self, id: cronoq::TaskId) -> Option<Bytes> {
90        if let Some(o) = self.outputs.get(id) {
91            return Some(o);
92        }
93        let q = self.queries.as_ref()?;
94        q.get_task(&cronoq::TaskRef::Id(id))
95            .await
96            .ok()
97            .flatten()
98            .and_then(|r| r.output)
99    }
100
101    /// Pool of a task, from the info cache or the database.
102    pub async fn pool_of(&self, id: cronoq::TaskId) -> Option<Arc<str>> {
103        if let Some(m) = self.info.get(id.0) {
104            return Some(m.pool);
105        }
106        let q = self.queries.as_ref()?;
107        q.get_task(&cronoq::TaskRef::Id(id))
108            .await
109            .ok()
110            .flatten()
111            .map(|r| Arc::from(r.pool.as_str()))
112    }
113
114    /// The pool's settings.
115    pub fn pool(&self, name: &str) -> Option<&PoolInfo> {
116        self.pools.get(name)
117    }
118}
cronoq_server/state.rs

cronoq_server/
state.rs
