ironflow_api/
worker_registry.rs

1//! In-memory record of the workers that recently asked for a run.
2//!
3//! Every call to the internal pick route carrying a `worker_id` records the
4//! capabilities the worker advertised. The run detail route reads it back to
5//! tell whether a pending run can be taken by any worker seen recently, so a
6//! run whose required tags no worker carries is flagged instead of waiting
7//! in silence.
8//!
9//! The record lives in the API process: with several API replicas, each one
10//! only knows the workers that polled it.
11
12use std::collections::HashMap;
13use std::sync::{Arc, PoisonError, RwLock};
14use std::time::{Duration, Instant};
15
16use ironflow_store::entities::WorkerCapabilities;
17
18use crate::entities::WorkerRouting;
19
20/// How long a worker counts as seen after its last pick request.
21///
22/// Workers poll every few seconds while idle; a worker busy on long runs
23/// still polls whenever a slot frees up.
24pub const WORKER_SEEN_TTL: Duration = Duration::from_secs(5 * 60);
25
26/// A worker as last seen by the pick route.
27#[derive(Debug, Clone)]
28struct SeenWorker {
29    /// `None` for a worker that predates routing: it takes every run.
30    capabilities: Option<WorkerCapabilities>,
31    last_seen: Instant,
32}
33
34/// Workers seen by the pick route within [`WORKER_SEEN_TTL`].
35///
36/// Cheap to clone: clones share the same record.
37///
38/// # Examples
39///
40/// ```
41/// use ironflow_api::worker_registry::WorkerRegistry;
42/// use ironflow_store::entities::WorkerCapabilities;
43///
44/// let registry = WorkerRegistry::new();
45/// registry.record("worker-1", Some(WorkerCapabilities::new(None, vec!["gpu".to_string()])));
46///
47/// let routing = registry.routing_for("transcode", &["gpu".to_string()]);
48/// assert_eq!(routing.seen_workers, 1);
49/// assert_eq!(routing.eligible_workers, 1);
50/// ```
51#[derive(Debug, Clone)]
52pub struct WorkerRegistry {
53    seen: Arc<RwLock<HashMap<String, SeenWorker>>>,
54    ttl: Duration,
55}
56
57impl Default for WorkerRegistry {
58    fn default() -> Self {
59        Self::with_ttl(WORKER_SEEN_TTL)
60    }
61}
62
63impl WorkerRegistry {
64    /// Create an empty registry using [`WORKER_SEEN_TTL`].
65    ///
66    /// # Examples
67    ///
68    /// ```
69    /// use ironflow_api::worker_registry::WorkerRegistry;
70    ///
71    /// let registry = WorkerRegistry::new();
72    /// assert_eq!(registry.routing_for("deploy", &[]).seen_workers, 0);
73    /// ```
74    pub fn new() -> Self {
75        Self::default()
76    }
77
78    /// Create an empty registry where a worker counts as seen for `ttl`.
79    ///
80    /// # Examples
81    ///
82    /// ```
83    /// use std::time::Duration;
84    ///
85    /// use ironflow_api::worker_registry::WorkerRegistry;
86    ///
87    /// let registry = WorkerRegistry::with_ttl(Duration::from_secs(30));
88    /// registry.record("worker-1", None);
89    /// assert_eq!(registry.routing_for("deploy", &[]).seen_workers, 1);
90    /// ```
91    pub fn with_ttl(ttl: Duration) -> Self {
92        Self {
93            seen: Arc::new(RwLock::new(HashMap::new())),
94            ttl,
95        }
96    }
97
98    /// Record that `worker_id` asked for a run with these capabilities.
99    ///
100    /// Replaces what was known about that worker, and drops the workers not
101    /// seen within the TTL.
102    ///
103    /// # Examples
104    ///
105    /// ```
106    /// use ironflow_api::worker_registry::WorkerRegistry;
107    ///
108    /// let registry = WorkerRegistry::new();
109    /// registry.record("worker-1", None);
110    /// registry.record("worker-1", None);
111    /// assert_eq!(registry.routing_for("deploy", &[]).seen_workers, 1);
112    /// ```
113    pub fn record(&self, worker_id: &str, capabilities: Option<WorkerCapabilities>) {
114        let now = Instant::now();
115        // The record is a cache: a writer that panicked leaves nothing worth
116        // refusing to read.
117        let mut workers = self.seen.write().unwrap_or_else(PoisonError::into_inner);
118        workers.retain(|_, worker| now.duration_since(worker.last_seen) < self.ttl);
119        workers.insert(
120            worker_id.to_string(),
121            SeenWorker {
122                capabilities,
123                last_seen: now,
124            },
125        );
126    }
127
128    /// Count the workers seen within the TTL, and those that could take a run
129    /// of `workflow_name` requiring `required_tags`.
130    ///
131    /// A worker without capabilities (an older release) can take any run.
132    ///
133    /// # Examples
134    ///
135    /// ```
136    /// use ironflow_api::worker_registry::WorkerRegistry;
137    /// use ironflow_store::entities::WorkerCapabilities;
138    ///
139    /// let registry = WorkerRegistry::new();
140    /// registry.record("cpu-1", Some(WorkerCapabilities::new(None, Vec::new())));
141    ///
142    /// let routing = registry.routing_for("transcode", &["gpu".to_string()]);
143    /// assert_eq!(routing.seen_workers, 1);
144    /// assert_eq!(routing.eligible_workers, 0);
145    /// ```
146    pub fn routing_for(&self, workflow_name: &str, required_tags: &[String]) -> WorkerRouting {
147        let now = Instant::now();
148        let workers = self.seen.read().unwrap_or_else(PoisonError::into_inner);
149        let seen: Vec<&SeenWorker> = workers
150            .values()
151            .filter(|seen| now.duration_since(seen.last_seen) < self.ttl)
152            .collect();
153        let eligible = seen
154            .iter()
155            .filter(|seen| {
156                seen.capabilities
157                    .as_ref()
158                    .is_none_or(|caps| caps.can_take(workflow_name, required_tags))
159            })
160            .count();
161
162        WorkerRouting {
163            seen_workers: saturating_u32(seen.len()),
164            eligible_workers: saturating_u32(eligible),
165        }
166    }
167}
168
169fn saturating_u32(count: usize) -> u32 {
170    u32::try_from(count).unwrap_or(u32::MAX)
171}
172
173#[cfg(test)]
174mod tests {
175    use std::thread::sleep;
176
177    use super::*;
178
179    fn strings(values: &[&str]) -> Vec<String> {
180        values.iter().map(|v| (*v).to_string()).collect()
181    }
182
183    fn caps(workflows: Option<&[&str]>, tags: &[&str]) -> Option<WorkerCapabilities> {
184        Some(WorkerCapabilities::new(
185            workflows.map(strings),
186            strings(tags),
187        ))
188    }
189
190    #[test]
191    fn empty_registry_has_no_worker() {
192        let routing = WorkerRegistry::new().routing_for("deploy", &[]);
193        assert_eq!(
194            routing,
195            WorkerRouting {
196                seen_workers: 0,
197                eligible_workers: 0
198            }
199        );
200    }
201
202    #[test]
203    fn worker_without_required_tag_is_seen_but_not_eligible() {
204        let registry = WorkerRegistry::new();
205        registry.record("cpu-1", caps(None, &["arm"]));
206        registry.record("gpu-1", caps(None, &["gpu"]));
207
208        let routing = registry.routing_for("transcode", &strings(&["gpu"]));
209        assert_eq!(routing.seen_workers, 2);
210        assert_eq!(routing.eligible_workers, 1);
211    }
212
213    #[test]
214    fn worker_without_the_workflow_is_not_eligible() {
215        let registry = WorkerRegistry::new();
216        registry.record("worker-1", caps(Some(&["deploy"]), &[]));
217
218        assert_eq!(registry.routing_for("deploy", &[]).eligible_workers, 1);
219        assert_eq!(registry.routing_for("build", &[]).eligible_workers, 0);
220    }
221
222    #[test]
223    fn legacy_worker_is_eligible_for_everything() {
224        let registry = WorkerRegistry::new();
225        registry.record("legacy", None);
226
227        let routing = registry.routing_for("transcode", &strings(&["gpu"]));
228        assert_eq!(routing.seen_workers, 1);
229        assert_eq!(routing.eligible_workers, 1);
230    }
231
232    #[test]
233    fn record_replaces_previous_capabilities() {
234        let registry = WorkerRegistry::new();
235        registry.record("worker-1", caps(None, &["gpu"]));
236        registry.record("worker-1", caps(None, &[]));
237
238        let routing = registry.routing_for("transcode", &strings(&["gpu"]));
239        assert_eq!(routing.seen_workers, 1);
240        assert_eq!(routing.eligible_workers, 0);
241    }
242
243    #[test]
244    fn workers_not_seen_within_ttl_are_ignored_and_pruned() {
245        let registry = WorkerRegistry::with_ttl(Duration::from_millis(20));
246        registry.record("old", None);
247        sleep(Duration::from_millis(40));
248
249        assert_eq!(registry.routing_for("deploy", &[]).seen_workers, 0);
250
251        registry.record("new", None);
252        let workers = registry.seen.read().unwrap();
253        assert_eq!(workers.len(), 1);
254        assert!(workers.contains_key("new"));
255    }
256
257    #[test]
258    fn clones_share_the_record() {
259        let registry = WorkerRegistry::new();
260        let clone = registry.clone();
261        clone.record("worker-1", None);
262        assert_eq!(registry.routing_for("deploy", &[]).seen_workers, 1);
263    }
264}
ironflow_api/worker_registry.rs

ironflow_api/
worker_registry.rs
