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