1#![allow(dead_code)]
2use std::collections::BinaryHeap;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
12pub enum PriorityLevel {
13 Low = 0,
15 Normal = 1,
17 High = 2,
19 Critical = 3,
21}
22
23impl PriorityLevel {
24 #[must_use]
26 pub fn weight(self) -> u32 {
27 self as u32
28 }
29
30 #[must_use]
32 pub fn promote(self) -> Self {
33 match self {
34 Self::Low => Self::Normal,
35 Self::Normal => Self::High,
36 Self::High | Self::Critical => Self::Critical,
37 }
38 }
39}
40
41impl std::fmt::Display for PriorityLevel {
42 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43 match self {
44 Self::Low => write!(f, "low"),
45 Self::Normal => write!(f, "normal"),
46 Self::High => write!(f, "high"),
47 Self::Critical => write!(f, "critical"),
48 }
49 }
50}
51
52#[derive(Debug, Clone)]
54pub struct PriorityEntry {
55 pub task_id: String,
57 pub priority: PriorityLevel,
59 pub base_priority: PriorityLevel,
61 pub deadline_secs: u64,
63 pub submitted_secs: u64,
65 insertion_order: u64,
67 pub age_count: u32,
69}
70
71impl PriorityEntry {
72 pub fn new(
74 task_id: impl Into<String>,
75 priority: PriorityLevel,
76 submitted_secs: u64,
77 insertion_order: u64,
78 ) -> Self {
79 Self {
80 task_id: task_id.into(),
81 priority,
82 base_priority: priority,
83 deadline_secs: 0,
84 submitted_secs,
85 insertion_order,
86 age_count: 0,
87 }
88 }
89
90 #[must_use]
92 pub fn with_deadline(mut self, deadline_secs: u64) -> Self {
93 self.deadline_secs = deadline_secs;
94 self
95 }
96
97 #[must_use]
99 pub fn has_deadline(&self) -> bool {
100 self.deadline_secs > 0
101 }
102
103 #[must_use]
105 pub fn is_overdue(&self, now_secs: u64) -> bool {
106 self.has_deadline() && now_secs > self.deadline_secs
107 }
108
109 #[must_use]
111 pub fn wait_time(&self, now_secs: u64) -> u64 {
112 now_secs.saturating_sub(self.submitted_secs)
113 }
114}
115
116impl PartialEq for PriorityEntry {
118 fn eq(&self, other: &Self) -> bool {
119 self.task_id == other.task_id
120 }
121}
122
123impl Eq for PriorityEntry {}
124
125impl PartialOrd for PriorityEntry {
126 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
127 Some(self.cmp(other))
128 }
129}
130
131impl Ord for PriorityEntry {
132 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
133 self.priority
135 .cmp(&other.priority)
136 .then_with(|| {
138 if self.deadline_secs == 0 && other.deadline_secs == 0 {
139 std::cmp::Ordering::Equal
140 } else if self.deadline_secs == 0 {
141 std::cmp::Ordering::Less } else if other.deadline_secs == 0 {
143 std::cmp::Ordering::Greater
144 } else {
145 other.deadline_secs.cmp(&self.deadline_secs) }
147 })
148 .then_with(|| other.submitted_secs.cmp(&self.submitted_secs))
150 .then_with(|| other.insertion_order.cmp(&self.insertion_order))
152 }
153}
154
155#[derive(Debug)]
157pub struct TaskPriorityQueue {
158 heap: BinaryHeap<PriorityEntry>,
160 insertion_counter: u64,
162 aging_threshold_secs: u64,
164}
165
166impl TaskPriorityQueue {
167 #[must_use]
169 pub fn new() -> Self {
170 Self {
171 heap: BinaryHeap::new(),
172 insertion_counter: 0,
173 aging_threshold_secs: 300, }
175 }
176
177 #[must_use]
179 pub fn with_aging_threshold(mut self, secs: u64) -> Self {
180 self.aging_threshold_secs = secs;
181 self
182 }
183
184 pub fn enqueue(
186 &mut self,
187 task_id: impl Into<String>,
188 priority: PriorityLevel,
189 submitted_secs: u64,
190 ) -> u64 {
191 let order = self.insertion_counter;
192 self.insertion_counter += 1;
193 let entry = PriorityEntry::new(task_id, priority, submitted_secs, order);
194 self.heap.push(entry);
195 order
196 }
197
198 pub fn enqueue_with_deadline(
200 &mut self,
201 task_id: impl Into<String>,
202 priority: PriorityLevel,
203 submitted_secs: u64,
204 deadline_secs: u64,
205 ) -> u64 {
206 let order = self.insertion_counter;
207 self.insertion_counter += 1;
208 let entry = PriorityEntry::new(task_id, priority, submitted_secs, order)
209 .with_deadline(deadline_secs);
210 self.heap.push(entry);
211 order
212 }
213
214 pub fn dequeue(&mut self) -> Option<PriorityEntry> {
216 self.heap.pop()
217 }
218
219 #[must_use]
221 pub fn peek(&self) -> Option<&PriorityEntry> {
222 self.heap.peek()
223 }
224
225 #[must_use]
227 pub fn len(&self) -> usize {
228 self.heap.len()
229 }
230
231 #[must_use]
233 pub fn is_empty(&self) -> bool {
234 self.heap.is_empty()
235 }
236
237 pub fn apply_aging(&mut self, now_secs: u64) {
241 let threshold = self.aging_threshold_secs;
242 let entries: Vec<PriorityEntry> = self.heap.drain().collect();
243 for mut entry in entries {
244 if entry.wait_time(now_secs) > threshold * (u64::from(entry.age_count) + 1) {
245 entry.priority = entry.priority.promote();
246 entry.age_count += 1;
247 }
248 self.heap.push(entry);
249 }
250 }
251
252 pub fn drain_overdue(&mut self, now_secs: u64) -> Vec<PriorityEntry> {
254 let mut overdue = Vec::new();
255 let mut remaining = Vec::new();
256 while let Some(entry) = self.heap.pop() {
257 if entry.is_overdue(now_secs) {
258 overdue.push(entry);
259 } else {
260 remaining.push(entry);
261 }
262 }
263 for e in remaining {
264 self.heap.push(e);
265 }
266 overdue
267 }
268
269 pub fn clear(&mut self) {
271 self.heap.clear();
272 }
273
274 #[must_use]
276 pub fn count_by_priority(&self) -> [usize; 4] {
277 let mut counts = [0_usize; 4];
278 for entry in &self.heap {
279 let idx = entry.priority.weight() as usize;
280 if idx < 4 {
281 counts[idx] += 1;
282 }
283 }
284 counts
285 }
286}
287
288impl Default for TaskPriorityQueue {
289 fn default() -> Self {
290 Self::new()
291 }
292}
293
294#[cfg(test)]
295mod tests {
296 use super::*;
297
298 #[test]
299 fn test_priority_level_ordering() {
300 assert!(PriorityLevel::Critical > PriorityLevel::High);
301 assert!(PriorityLevel::High > PriorityLevel::Normal);
302 assert!(PriorityLevel::Normal > PriorityLevel::Low);
303 }
304
305 #[test]
306 fn test_priority_level_promote() {
307 assert_eq!(PriorityLevel::Low.promote(), PriorityLevel::Normal);
308 assert_eq!(PriorityLevel::Normal.promote(), PriorityLevel::High);
309 assert_eq!(PriorityLevel::High.promote(), PriorityLevel::Critical);
310 assert_eq!(PriorityLevel::Critical.promote(), PriorityLevel::Critical);
311 }
312
313 #[test]
314 fn test_priority_level_display() {
315 assert_eq!(format!("{}", PriorityLevel::Low), "low");
316 assert_eq!(format!("{}", PriorityLevel::Normal), "normal");
317 assert_eq!(format!("{}", PriorityLevel::High), "high");
318 assert_eq!(format!("{}", PriorityLevel::Critical), "critical");
319 }
320
321 #[test]
322 fn test_priority_entry_deadline() {
323 let entry = PriorityEntry::new("t1", PriorityLevel::Normal, 1000, 0).with_deadline(2000);
324 assert!(entry.has_deadline());
325 assert!(!entry.is_overdue(1500));
326 assert!(entry.is_overdue(2500));
327 }
328
329 #[test]
330 fn test_priority_entry_no_deadline() {
331 let entry = PriorityEntry::new("t1", PriorityLevel::Normal, 1000, 0);
332 assert!(!entry.has_deadline());
333 assert!(!entry.is_overdue(9999));
334 }
335
336 #[test]
337 fn test_priority_entry_wait_time() {
338 let entry = PriorityEntry::new("t1", PriorityLevel::Normal, 1000, 0);
339 assert_eq!(entry.wait_time(1500), 500);
340 assert_eq!(entry.wait_time(500), 0); }
342
343 #[test]
344 fn test_queue_enqueue_dequeue_priority() {
345 let mut q = TaskPriorityQueue::new();
346 q.enqueue("low", PriorityLevel::Low, 1000);
347 q.enqueue("critical", PriorityLevel::Critical, 1000);
348 q.enqueue("normal", PriorityLevel::Normal, 1000);
349
350 let first = q.dequeue().expect("should succeed in test");
351 assert_eq!(first.task_id, "critical");
352 let second = q.dequeue().expect("should succeed in test");
353 assert_eq!(second.task_id, "normal");
354 let third = q.dequeue().expect("should succeed in test");
355 assert_eq!(third.task_id, "low");
356 }
357
358 #[test]
359 fn test_queue_fifo_within_same_priority() {
360 let mut q = TaskPriorityQueue::new();
361 q.enqueue("a", PriorityLevel::Normal, 1000);
362 q.enqueue("b", PriorityLevel::Normal, 1000);
363 q.enqueue("c", PriorityLevel::Normal, 1000);
364
365 assert_eq!(q.dequeue().expect("should succeed in test").task_id, "a");
366 assert_eq!(q.dequeue().expect("should succeed in test").task_id, "b");
367 assert_eq!(q.dequeue().expect("should succeed in test").task_id, "c");
368 }
369
370 #[test]
371 fn test_queue_deadline_ordering() {
372 let mut q = TaskPriorityQueue::new();
373 q.enqueue_with_deadline("far", PriorityLevel::Normal, 1000, 5000);
374 q.enqueue_with_deadline("near", PriorityLevel::Normal, 1000, 2000);
375
376 assert_eq!(q.dequeue().expect("should succeed in test").task_id, "near");
378 assert_eq!(q.dequeue().expect("should succeed in test").task_id, "far");
379 }
380
381 #[test]
382 fn test_queue_len_and_empty() {
383 let mut q = TaskPriorityQueue::new();
384 assert!(q.is_empty());
385 assert_eq!(q.len(), 0);
386 q.enqueue("a", PriorityLevel::Normal, 1000);
387 assert!(!q.is_empty());
388 assert_eq!(q.len(), 1);
389 }
390
391 #[test]
392 fn test_queue_peek() {
393 let mut q = TaskPriorityQueue::new();
394 assert!(q.peek().is_none());
395 q.enqueue("a", PriorityLevel::High, 1000);
396 q.enqueue("b", PriorityLevel::Low, 1000);
397 assert_eq!(q.peek().expect("should succeed in test").task_id, "a");
398 assert_eq!(q.len(), 2); }
400
401 #[test]
402 fn test_queue_apply_aging() {
403 let mut q = TaskPriorityQueue::new().with_aging_threshold(100);
404 q.enqueue("old", PriorityLevel::Low, 0);
405 q.enqueue("new", PriorityLevel::Low, 900);
406
407 q.apply_aging(200);
409
410 let first = q.dequeue().expect("should succeed in test");
411 assert_eq!(first.task_id, "old");
412 assert_eq!(first.priority, PriorityLevel::Normal);
413 assert_eq!(first.age_count, 1);
414 }
415
416 #[test]
417 fn test_queue_drain_overdue() {
418 let mut q = TaskPriorityQueue::new();
419 q.enqueue_with_deadline("overdue", PriorityLevel::Normal, 1000, 1500);
420 q.enqueue_with_deadline("ok", PriorityLevel::Normal, 1000, 3000);
421 q.enqueue("no-deadline", PriorityLevel::Normal, 1000);
422
423 let overdue = q.drain_overdue(2000);
424 assert_eq!(overdue.len(), 1);
425 assert_eq!(overdue[0].task_id, "overdue");
426 assert_eq!(q.len(), 2);
427 }
428
429 #[test]
430 fn test_queue_clear() {
431 let mut q = TaskPriorityQueue::new();
432 q.enqueue("a", PriorityLevel::Normal, 1000);
433 q.enqueue("b", PriorityLevel::High, 1000);
434 q.clear();
435 assert!(q.is_empty());
436 }
437
438 #[test]
439 fn test_queue_count_by_priority() {
440 let mut q = TaskPriorityQueue::new();
441 q.enqueue("a", PriorityLevel::Low, 0);
442 q.enqueue("b", PriorityLevel::Normal, 0);
443 q.enqueue("c", PriorityLevel::Normal, 0);
444 q.enqueue("d", PriorityLevel::Critical, 0);
445
446 let counts = q.count_by_priority();
447 assert_eq!(counts[PriorityLevel::Low.weight() as usize], 1);
448 assert_eq!(counts[PriorityLevel::Normal.weight() as usize], 2);
449 assert_eq!(counts[PriorityLevel::High.weight() as usize], 0);
450 assert_eq!(counts[PriorityLevel::Critical.weight() as usize], 1);
451 }
452}