aion_server/worker/queue_service/
state.rs

1//! The live, queryable unserved-queue state.
2//!
3//! A parked dispatch is no longer invisible. While one waits, its address sits
4//! in here with the taxonomy reason, the policy it is parked under, the poller
5//! census behind the verdict, and every run waiting on it — so "which runs are
6//! stuck, on what, since when" is a server-side read rather than a log
7//! archaeology exercise.
8//!
9//! Entries exist only while a dispatch is actually parked: the wait loop clears
10//! its own entry on every exit path (served, refused, drained), so the state
11//! never accumulates addresses that are fine.
12
13use std::collections::{BTreeMap, HashMap};
14use std::sync::{Arc, Mutex};
15use std::time::{Duration, Instant};
16
17use aion_core::{ActivityId, WorkflowId};
18
19use super::census::PoolCensus;
20use super::policy::QueueServicePolicy;
21use super::taxonomy::{QueueServiceReason, ServiceAddress};
22use crate::error::ServerError;
23
24/// Identity of an unserved address: the pool plus the activity type on it.
25///
26/// The node pin is per-dispatch, not per-address, so it lives on the waiting
27/// entries rather than the key.
28#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
29pub struct UnservedKey {
30    /// Correctness/isolation boundary.
31    pub namespace: String,
32    /// Pool selector within the namespace.
33    pub task_queue: String,
34    /// Activity type that is unserved on it.
35    pub activity_type: String,
36}
37
38impl UnservedKey {
39    /// The key an address belongs to.
40    #[must_use]
41    pub fn of(address: &ServiceAddress) -> Self {
42        Self {
43            namespace: address.namespace.clone(),
44            task_queue: address.task_queue.clone(),
45            activity_type: address.activity_type.clone(),
46        }
47    }
48}
49
50/// One run's activity waiting on an unserved address.
51#[derive(Clone, Debug, Eq, PartialEq)]
52pub struct UnservedDispatch {
53    /// Owning workflow.
54    pub workflow_id: WorkflowId,
55    /// Activity ordinal recorded in history.
56    pub activity_id: ActivityId,
57    /// Node this dispatch is pinned to, if any.
58    pub node: Option<String>,
59    /// How long this dispatch has been parked.
60    pub waiting_for: Duration,
61}
62
63/// The unserved state of one address, as read by an operator surface.
64#[derive(Clone, Debug, Eq, PartialEq)]
65pub struct UnservedQueue {
66    /// Which address is unserved.
67    pub key: UnservedKey,
68    /// Taxonomy state of the address at the last classification.
69    pub reason: QueueServiceReason,
70    /// Policy the parked dispatches are being held under.
71    pub policy: QueueServicePolicy,
72    /// Live poller census behind the verdict.
73    pub census: PoolCensus,
74    /// How long this address has been continuously unserved.
75    pub unserved_for: Duration,
76    /// Every dispatch currently parked on it, oldest first.
77    pub waiting: Vec<UnservedDispatch>,
78}
79
80#[derive(Clone, Debug)]
81struct Entry {
82    reason: QueueServiceReason,
83    policy: QueueServicePolicy,
84    census: PoolCensus,
85    since: Instant,
86    /// Keyed by execution identity; `WorkflowId`/`ActivityId` are hashable but
87    /// not ordered, so the read path sorts explicitly rather than relying on
88    /// map order.
89    waiting: HashMap<(WorkflowId, ActivityId), Waiting>,
90}
91
92#[derive(Clone, Debug)]
93struct Waiting {
94    node: Option<String>,
95    since: Instant,
96}
97
98/// Shared handle onto the live unserved-queue state.
99#[derive(Clone, Debug, Default)]
100pub struct QueueServiceState {
101    inner: Arc<Mutex<BTreeMap<UnservedKey, Entry>>>,
102}
103
104impl QueueServiceState {
105    /// Record (or refresh) one dispatch parked on an unserved address.
106    ///
107    /// # Errors
108    ///
109    /// Returns [`ServerError::LockPoisoned`] if the state lock is poisoned.
110    pub fn mark(&self, parked: Parked<'_>) -> Result<(), ServerError> {
111        let key = UnservedKey::of(parked.address);
112        let now = Instant::now();
113        let mut state = self.state()?;
114        let entry = state.entry(key).or_insert_with(|| Entry {
115            reason: parked.reason,
116            policy: parked.policy,
117            census: parked.census,
118            since: now,
119            waiting: HashMap::new(),
120        });
121        entry.reason = parked.reason;
122        entry.policy = parked.policy;
123        entry.census = parked.census;
124        entry
125            .waiting
126            .entry((parked.workflow_id.clone(), parked.activity_id.clone()))
127            .or_insert_with(|| Waiting {
128                node: parked.address.node.clone(),
129                since: now,
130            });
131        Ok(())
132    }
133
134    /// Drop one dispatch from the unserved state, dropping the address itself
135    /// when nothing waits on it any more.
136    ///
137    /// # Errors
138    ///
139    /// Returns [`ServerError::LockPoisoned`] if the state lock is poisoned.
140    pub fn clear(
141        &self,
142        address: &ServiceAddress,
143        workflow_id: &WorkflowId,
144        activity_id: &ActivityId,
145    ) -> Result<(), ServerError> {
146        let key = UnservedKey::of(address);
147        let mut state = self.state()?;
148        let Some(entry) = state.get_mut(&key) else {
149            return Ok(());
150        };
151        entry
152            .waiting
153            .remove(&(workflow_id.clone(), activity_id.clone()));
154        if entry.waiting.is_empty() {
155            state.remove(&key);
156        }
157        Ok(())
158    }
159
160    /// Every address currently unserved, in stable address order.
161    ///
162    /// # Errors
163    ///
164    /// Returns [`ServerError::LockPoisoned`] if the state lock is poisoned.
165    pub fn unserved(&self) -> Result<Vec<UnservedQueue>, ServerError> {
166        let now = Instant::now();
167        let state = self.state()?;
168        Ok(state
169            .iter()
170            .map(|(key, entry)| {
171                let mut waiting: Vec<UnservedDispatch> = entry
172                    .waiting
173                    .iter()
174                    .map(|((workflow_id, activity_id), held)| UnservedDispatch {
175                        workflow_id: workflow_id.clone(),
176                        activity_id: activity_id.clone(),
177                        node: held.node.clone(),
178                        waiting_for: now.saturating_duration_since(held.since),
179                    })
180                    .collect();
181                waiting.sort_by(|left, right| {
182                    right.waiting_for.cmp(&left.waiting_for).then_with(|| {
183                        left.workflow_id
184                            .to_string()
185                            .cmp(&right.workflow_id.to_string())
186                    })
187                });
188                UnservedQueue {
189                    key: key.clone(),
190                    reason: entry.reason,
191                    policy: entry.policy,
192                    census: entry.census,
193                    unserved_for: now.saturating_duration_since(entry.since),
194                    waiting,
195                }
196            })
197            .collect())
198    }
199
200    /// How many dispatches are currently parked across every address on
201    /// `task_queue`, in any namespace.
202    ///
203    /// Read by the worker-expiry sweep so the deregistration of a dead worker
204    /// states the operational consequence in the same breath as the cause: "this
205    /// worker is gone AND seven dispatches are already parked on the queue it
206    /// was serving" is an alertable sentence; "worker deregistered" alone is
207    /// not. Queue-scoped rather than address-scoped because a worker serves a
208    /// whole queue, not one activity type on it.
209    ///
210    /// # Errors
211    ///
212    /// Returns [`ServerError::LockPoisoned`] if the state lock is poisoned.
213    pub fn parked_on_queue(&self, task_queue: &str) -> Result<usize, ServerError> {
214        Ok(self
215            .state()?
216            .iter()
217            .filter(|(key, _)| key.task_queue == task_queue)
218            .map(|(_, entry)| entry.waiting.len())
219            .sum())
220    }
221
222    /// Whether one run's activity is currently parked on an unserved queue.
223    ///
224    /// # Errors
225    ///
226    /// Returns [`ServerError::LockPoisoned`] if the state lock is poisoned.
227    pub fn is_unserved(
228        &self,
229        workflow_id: &WorkflowId,
230        activity_id: &ActivityId,
231    ) -> Result<bool, ServerError> {
232        let held = (workflow_id.clone(), activity_id.clone());
233        Ok(self
234            .state()?
235            .values()
236            .any(|entry| entry.waiting.contains_key(&held)))
237    }
238
239    fn state(
240        &self,
241    ) -> Result<std::sync::MutexGuard<'_, BTreeMap<UnservedKey, Entry>>, ServerError> {
242        self.inner
243            .lock()
244            .map_err(|_| ServerError::lock_poisoned("queue service state"))
245    }
246}
247
248/// One parked dispatch, as handed to [`QueueServiceState::mark`].
249#[derive(Clone, Copy, Debug)]
250pub struct Parked<'a> {
251    /// Address the dispatch is parked on.
252    pub address: &'a ServiceAddress,
253    /// Taxonomy state of that address.
254    pub reason: QueueServiceReason,
255    /// Policy the dispatch is parked under.
256    pub policy: QueueServicePolicy,
257    /// Poller census behind the verdict.
258    pub census: PoolCensus,
259    /// Owning workflow.
260    pub workflow_id: &'a WorkflowId,
261    /// Activity ordinal.
262    pub activity_id: &'a ActivityId,
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    fn address() -> ServiceAddress {
270        ServiceAddress {
271            namespace: "default".to_owned(),
272            task_queue: "general".to_owned(),
273            activity_type: "greet".to_owned(),
274            node: None,
275        }
276    }
277
278    fn parked<'a>(
279        address: &'a ServiceAddress,
280        workflow_id: &'a WorkflowId,
281        activity_id: &'a ActivityId,
282    ) -> Parked<'a> {
283        Parked {
284            address,
285            reason: QueueServiceReason::NoLivePollers,
286            policy: QueueServicePolicy::DurablePending,
287            census: PoolCensus::default(),
288            workflow_id,
289            activity_id,
290        }
291    }
292
293    #[test]
294    fn a_parked_dispatch_is_visible_immediately() -> Result<(), ServerError> {
295        let state = QueueServiceState::default();
296        let address = address();
297        let workflow_id = WorkflowId::new_v4();
298        let activity_id = ActivityId::from_sequence_position(3);
299        state.mark(parked(&address, &workflow_id, &activity_id))?;
300
301        let unserved = state.unserved()?;
302        assert_eq!(unserved.len(), 1);
303        assert_eq!(unserved[0].reason, QueueServiceReason::NoLivePollers);
304        assert_eq!(unserved[0].policy, QueueServicePolicy::DurablePending);
305        assert_eq!(unserved[0].key.task_queue, "general");
306        assert_eq!(unserved[0].waiting.len(), 1);
307        assert_eq!(unserved[0].waiting[0].workflow_id, workflow_id);
308        assert!(state.is_unserved(&workflow_id, &activity_id)?);
309        Ok(())
310    }
311
312    #[test]
313    fn clearing_the_last_waiter_removes_the_address() -> Result<(), ServerError> {
314        let state = QueueServiceState::default();
315        let address = address();
316        let first = WorkflowId::new_v4();
317        let second = WorkflowId::new_v4();
318        let activity_id = ActivityId::from_sequence_position(0);
319        state.mark(parked(&address, &first, &activity_id))?;
320        state.mark(parked(&address, &second, &activity_id))?;
321        assert_eq!(state.unserved()?.len(), 1);
322        assert_eq!(state.unserved()?[0].waiting.len(), 2);
323
324        state.clear(&address, &first, &activity_id)?;
325        assert!(!state.is_unserved(&first, &activity_id)?);
326        assert_eq!(state.unserved()?[0].waiting.len(), 1);
327
328        state.clear(&address, &second, &activity_id)?;
329        assert!(state.unserved()?.is_empty(), "the address must disappear");
330        Ok(())
331    }
332
333    #[test]
334    fn clearing_an_unknown_dispatch_is_a_no_op() -> Result<(), ServerError> {
335        let state = QueueServiceState::default();
336        state.clear(
337            &address(),
338            &WorkflowId::new_v4(),
339            &ActivityId::from_sequence_position(0),
340        )?;
341        assert!(state.unserved()?.is_empty());
342        Ok(())
343    }
344
345    #[test]
346    fn re_marking_refreshes_the_verdict_without_duplicating_the_waiter() -> Result<(), ServerError>
347    {
348        let state = QueueServiceState::default();
349        let address = address();
350        let workflow_id = WorkflowId::new_v4();
351        let activity_id = ActivityId::from_sequence_position(0);
352        state.mark(parked(&address, &workflow_id, &activity_id))?;
353        state.mark(Parked {
354            reason: QueueServiceReason::PollersIncompatible,
355            census: PoolCensus {
356                workers_in_pool: 2,
357                ..PoolCensus::default()
358            },
359            ..parked(&address, &workflow_id, &activity_id)
360        })?;
361
362        let unserved = state.unserved()?;
363        assert_eq!(unserved.len(), 1);
364        assert_eq!(unserved[0].waiting.len(), 1);
365        assert_eq!(unserved[0].reason, QueueServiceReason::PollersIncompatible);
366        assert_eq!(unserved[0].census.workers_in_pool, 2);
367        Ok(())
368    }
369}
aion_server/worker/queue_service/state.rs

aion_server/worker/queue_service/
state.rs
