orch/engine/
priority.rs

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

orch/engine/
priority.rs
