orch/engine/
priority.rs

1use std::str::FromStr;
2use std::sync::{Arc, OnceLock};
3
4use tokio::sync::{OwnedSemaphorePermit, Semaphore};
5
6const HIGH_LABEL: &str = "priority:high";
7const LOW_LABEL: &str = "priority:low";
8const PREFIX: &str = "priority:";
9
10/// Dispatch priority read from the `priority:high` and `priority:low` labels
11#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
12pub enum Priority {
13    Low,
14    #[default]
15    Normal,
16    High,
17    /// A task with a linked PR, ranked above `priority:high` to ship pushed work first
18    OpenPr,
19}
20
21impl Priority {
22    /// The highest label wins when a task carries both
23    pub fn from_labels(labels: &[String]) -> Self {
24        let mut found = Self::Normal;
25        for label in labels {
26            match label.as_str() {
27                HIGH_LABEL => return Self::High,
28                LOW_LABEL => found = Self::Low,
29                _ => {}
30            }
31        }
32        found
33    }
34
35    pub fn as_str(self) -> &'static str {
36        match self {
37            Self::Low => "low",
38            Self::Normal => "normal",
39            Self::High => "high",
40            Self::OpenPr => "open-pr",
41        }
42    }
43
44    /// `OpenPr` when the task has a linked PR, else the label priority
45    #[must_use]
46    pub fn with_pr(self, has_pr: bool) -> Self {
47        if has_pr { Self::OpenPr } else { self }
48    }
49
50    fn label(self) -> Option<&'static str> {
51        match self {
52            Self::Low => Some(LOW_LABEL),
53            Self::Normal | Self::OpenPr => None,
54            Self::High => Some(HIGH_LABEL),
55        }
56    }
57
58    /// `labels` with every priority label replaced by the one for `self`
59    pub fn apply_to(self, labels: &[String]) -> Vec<String> {
60        let mut out: Vec<String> = labels
61            .iter()
62            .filter(|l| !is_priority_label(l))
63            .cloned()
64            .collect();
65        if let Some(label) = self.label() {
66            out.push(label.to_string());
67        }
68        out
69    }
70}
71
72impl FromStr for Priority {
73    type Err = anyhow::Error;
74
75    fn from_str(s: &str) -> Result<Self, Self::Err> {
76        match s.trim().to_ascii_lowercase().as_str() {
77            "high" => Ok(Self::High),
78            "normal" => Ok(Self::Normal),
79            "low" => Ok(Self::Low),
80            other => anyhow::bail!("unknown priority '{other}', use high, normal or low"),
81        }
82    }
83}
84
85pub fn is_priority_label(label: &str) -> bool {
86    label.starts_with(PREFIX)
87}
88
89/// Highest priority first, arrival order kept inside a priority
90pub fn sort_by_priority<T>(tasks: &mut [T], priority_of: impl Fn(&T) -> Priority) {
91    tasks.sort_by_key(|t| std::cmp::Reverse(priority_of(t)));
92}
93
94/// The one slot a `priority:high` or open PR task may take beyond `max_parallel`
95fn extra_slot() -> &'static Arc<Semaphore> {
96    static SLOT: OnceLock<Arc<Semaphore>> = OnceLock::new();
97    SLOT.get_or_init(|| Arc::new(Semaphore::new(1)))
98}
99
100/// A regular slot, or for a high priority task the extra slot when every regular slot is busy
101pub fn acquire_dispatch_slot(
102    semaphore: &Arc<Semaphore>,
103    priority: Priority,
104) -> Option<OwnedSemaphorePermit> {
105    acquire_with_extra(semaphore, extra_slot(), priority)
106}
107
108fn acquire_with_extra(
109    semaphore: &Arc<Semaphore>,
110    extra: &Arc<Semaphore>,
111    priority: Priority,
112) -> Option<OwnedSemaphorePermit> {
113    if let Ok(permit) = Arc::clone(semaphore).try_acquire_owned() {
114        return Some(permit);
115    }
116
117    if priority >= Priority::High {
118        return Arc::clone(extra).try_acquire_owned().ok();
119    }
120
121    None
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    fn labels(l: &[&str]) -> Vec<String> {
129        l.iter().map(|s| (*s).to_string()).collect()
130    }
131
132    #[test]
133    fn reads_priority_from_labels() {
134        assert_eq!(Priority::from_labels(&labels(&["bug"])), Priority::Normal);
135        assert_eq!(
136            Priority::from_labels(&labels(&["priority:low"])),
137            Priority::Low
138        );
139        assert_eq!(
140            Priority::from_labels(&labels(&["priority:low", "priority:high"])),
141            Priority::High
142        );
143    }
144
145    #[test]
146    fn high_task_filed_later_dispatches_before_normal_task() {
147        let mut tasks = vec![
148            ("normal-first", labels(&["bug"])),
149            ("low", labels(&["priority:low"])),
150            ("high-later", labels(&["priority:high"])),
151            ("normal-second", labels(&[])),
152        ];
153        sort_by_priority(&mut tasks, |t| Priority::from_labels(&t.1));
154        let order: Vec<_> = tasks.iter().map(|t| t.0).collect();
155        assert_eq!(
156            order,
157            ["high-later", "normal-first", "normal-second", "low"]
158        );
159    }
160
161    #[test]
162    fn open_pr_task_dispatches_before_high_task() {
163        let mut tasks = vec![
164            ("high", Priority::High.with_pr(false)),
165            ("normal-with-pr", Priority::Normal.with_pr(true)),
166            ("normal", Priority::Normal.with_pr(false)),
167            ("low-with-pr", Priority::Low.with_pr(true)),
168        ];
169        sort_by_priority(&mut tasks, |t| t.1);
170        let order: Vec<_> = tasks.iter().map(|t| t.0).collect();
171        assert_eq!(order, ["normal-with-pr", "low-with-pr", "high", "normal"]);
172    }
173
174    #[test]
175    fn open_pr_task_may_use_the_extra_slot() {
176        let semaphore = Arc::new(Semaphore::new(1));
177        let extra = Arc::new(Semaphore::new(1));
178        let _busy = acquire_with_extra(&semaphore, &extra, Priority::Normal).unwrap();
179        assert!(acquire_with_extra(&semaphore, &extra, Priority::OpenPr).is_some());
180    }
181
182    #[test]
183    fn high_task_takes_the_extra_slot_and_a_second_one_waits() {
184        let semaphore = Arc::new(Semaphore::new(1));
185        let extra = Arc::new(Semaphore::new(1));
186        let _busy = acquire_with_extra(&semaphore, &extra, Priority::Normal).unwrap();
187
188        assert!(acquire_with_extra(&semaphore, &extra, Priority::Normal).is_none());
189        assert!(acquire_with_extra(&semaphore, &extra, Priority::Low).is_none());
190
191        let first_high = acquire_with_extra(&semaphore, &extra, Priority::High);
192        assert!(first_high.is_some());
193        assert!(acquire_with_extra(&semaphore, &extra, Priority::High).is_none());
194
195        drop(first_high);
196        assert!(acquire_with_extra(&semaphore, &extra, Priority::High).is_some());
197    }
198
199    #[test]
200    fn high_task_prefers_a_regular_slot() {
201        let semaphore = Arc::new(Semaphore::new(1));
202        let extra = Arc::new(Semaphore::new(1));
203        let _permit = acquire_with_extra(&semaphore, &extra, Priority::High).unwrap();
204        assert_eq!(extra.available_permits(), 1);
205    }
206
207    #[test]
208    fn apply_replaces_existing_priority_labels() {
209        let current = labels(&["bug", "priority:low"]);
210        assert_eq!(
211            Priority::High.apply_to(&current),
212            labels(&["bug", "priority:high"])
213        );
214        assert_eq!(Priority::Normal.apply_to(&current), labels(&["bug"]));
215    }
216
217    #[test]
218    fn parses_priority_names() {
219        assert_eq!("High".parse::<Priority>().unwrap(), Priority::High);
220        assert!("urgent".parse::<Priority>().is_err());
221    }
222}
orch/engine/priority.rs

orch/engine/
priority.rs
