1#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
16pub enum IOPriority {
17 Background = 0,
19 Normal = 1,
21 High = 2,
23 Realtime = 3,
25}
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
33pub enum IODirection {
34 Read,
36 Write,
38}
39
40#[derive(Debug, Clone)]
46pub struct IORequest {
47 pub id: u64,
49 pub block_cid: String,
51 pub direction: IODirection,
53 pub priority: IOPriority,
55 pub size_bytes: u64,
57 pub deadline_tick: Option<u64>,
59 pub enqueued_tick: u64,
61}
62
63#[derive(Debug, Clone)]
69pub struct SchedulerConfig {
70 pub max_pending: usize,
72 pub read_weight: f64,
74 pub write_weight: f64,
76}
77
78impl Default for SchedulerConfig {
79 fn default() -> Self {
80 Self {
81 max_pending: 1000,
82 read_weight: 0.7,
83 write_weight: 0.3,
84 }
85 }
86}
87
88#[derive(Debug, Clone)]
94pub struct IOSchedulerStats {
95 pub pending_count: usize,
97 pub completed_reads: u64,
99 pub completed_writes: u64,
101 pub total_bytes_scheduled: u64,
103 pub expired_count: usize,
105}
106
107pub struct StorageIOScheduler {
120 config: SchedulerConfig,
121 pending: Vec<IORequest>,
122 next_id: u64,
123 current_tick: u64,
124 completed_reads: u64,
125 completed_writes: u64,
126 total_bytes_scheduled: u64,
127}
128
129impl StorageIOScheduler {
130 pub fn new(config: SchedulerConfig) -> Self {
132 Self {
133 config,
134 pending: Vec::new(),
135 next_id: 0,
136 current_tick: 0,
137 completed_reads: 0,
138 completed_writes: 0,
139 total_bytes_scheduled: 0,
140 }
141 }
142
143 pub fn submit(
148 &mut self,
149 block_cid: &str,
150 direction: IODirection,
151 priority: IOPriority,
152 size_bytes: u64,
153 deadline_tick: Option<u64>,
154 ) -> Result<u64, String> {
155 if self.pending.len() >= self.config.max_pending {
156 return Err(format!(
157 "pending queue full ({} / {})",
158 self.pending.len(),
159 self.config.max_pending,
160 ));
161 }
162
163 let id = self.next_id;
164 self.next_id = self.next_id.saturating_add(1);
165
166 self.pending.push(IORequest {
167 id,
168 block_cid: block_cid.to_string(),
169 direction,
170 priority,
171 size_bytes,
172 deadline_tick,
173 enqueued_tick: self.current_tick,
174 });
175
176 Ok(id)
177 }
178
179 pub fn next_request(&mut self) -> Option<IORequest> {
186 if self.pending.is_empty() {
187 return None;
188 }
189
190 let tick = self.current_tick;
191
192 let mut best_idx: usize = 0;
194 let mut best_urgent = false;
195 let mut best_deadline: u64 = u64::MAX;
196 let mut best_priority = IOPriority::Background;
197 let mut best_enqueued: u64 = u64::MAX;
198
199 for (i, req) in self.pending.iter().enumerate() {
200 let urgent = req.deadline_tick.map(|d| d <= tick).unwrap_or(false);
201
202 let better = if urgent && !best_urgent {
203 true
205 } else if urgent && best_urgent {
206 let dl = req.deadline_tick.unwrap_or(u64::MAX);
208 dl < best_deadline || (dl == best_deadline && req.enqueued_tick < best_enqueued)
209 } else if !urgent && best_urgent {
210 false
211 } else {
212 req.priority > best_priority
214 || (req.priority == best_priority && req.enqueued_tick < best_enqueued)
215 };
216
217 if better {
218 best_idx = i;
219 best_urgent = urgent;
220 best_deadline = req.deadline_tick.unwrap_or(u64::MAX);
221 best_priority = req.priority;
222 best_enqueued = req.enqueued_tick;
223 }
224 }
225
226 Some(self.pending.remove(best_idx))
227 }
228
229 pub fn complete(&mut self, request_id: u64, direction: IODirection, size_bytes: u64) {
233 match direction {
234 IODirection::Read => self.completed_reads = self.completed_reads.saturating_add(1),
235 IODirection::Write => self.completed_writes = self.completed_writes.saturating_add(1),
236 }
237 self.total_bytes_scheduled = self.total_bytes_scheduled.saturating_add(size_bytes);
238 self.pending.retain(|r| r.id != request_id);
240 }
241
242 pub fn cancel(&mut self, request_id: u64) -> bool {
246 let before = self.pending.len();
247 self.pending.retain(|r| r.id != request_id);
248 self.pending.len() < before
249 }
250
251 pub fn pending_count(&self) -> usize {
253 self.pending.len()
254 }
255
256 pub fn pending_reads(&self) -> usize {
258 self.pending
259 .iter()
260 .filter(|r| r.direction == IODirection::Read)
261 .count()
262 }
263
264 pub fn pending_writes(&self) -> usize {
266 self.pending
267 .iter()
268 .filter(|r| r.direction == IODirection::Write)
269 .count()
270 }
271
272 pub fn tick(&mut self) {
274 self.current_tick = self.current_tick.saturating_add(1);
275 }
276
277 pub fn current_tick(&self) -> u64 {
279 self.current_tick
280 }
281
282 pub fn expired_requests(&self) -> Vec<&IORequest> {
284 let tick = self.current_tick;
285 self.pending
286 .iter()
287 .filter(|r| r.deadline_tick.map(|d| d < tick).unwrap_or(false))
288 .collect()
289 }
290
291 pub fn drain_expired(&mut self) -> Vec<IORequest> {
293 let tick = self.current_tick;
294 let mut expired = Vec::new();
295 let mut kept = Vec::new();
296 for req in self.pending.drain(..) {
297 if req.deadline_tick.map(|d| d < tick).unwrap_or(false) {
298 expired.push(req);
299 } else {
300 kept.push(req);
301 }
302 }
303 self.pending = kept;
304 expired
305 }
306
307 pub fn stats(&self) -> IOSchedulerStats {
309 IOSchedulerStats {
310 pending_count: self.pending.len(),
311 completed_reads: self.completed_reads,
312 completed_writes: self.completed_writes,
313 total_bytes_scheduled: self.total_bytes_scheduled,
314 expired_count: self.expired_requests().len(),
315 }
316 }
317
318 pub fn config(&self) -> &SchedulerConfig {
320 &self.config
321 }
322}
323
324#[cfg(test)]
329mod tests {
330 use super::*;
331
332 fn default_scheduler() -> StorageIOScheduler {
333 StorageIOScheduler::new(SchedulerConfig::default())
334 }
335
336 fn small_scheduler(max: usize) -> StorageIOScheduler {
337 StorageIOScheduler::new(SchedulerConfig {
338 max_pending: max,
339 ..SchedulerConfig::default()
340 })
341 }
342
343 #[test]
346 fn test_submit_returns_unique_ids() {
347 let mut s = default_scheduler();
348 let id0 = s.submit("cid0", IODirection::Read, IOPriority::Normal, 100, None);
349 let id1 = s.submit("cid1", IODirection::Write, IOPriority::Normal, 200, None);
350 assert!(id0.is_ok());
351 assert!(id1.is_ok());
352 assert_ne!(id0.ok(), id1.ok());
353 }
354
355 #[test]
356 fn test_submit_increments_pending() {
357 let mut s = default_scheduler();
358 assert_eq!(s.pending_count(), 0);
359 let _ = s.submit("c", IODirection::Read, IOPriority::Normal, 10, None);
360 assert_eq!(s.pending_count(), 1);
361 let _ = s.submit("c", IODirection::Write, IOPriority::High, 20, None);
362 assert_eq!(s.pending_count(), 2);
363 }
364
365 #[test]
366 fn test_submit_max_pending_enforced() {
367 let mut s = small_scheduler(2);
368 assert!(s
369 .submit("a", IODirection::Read, IOPriority::Normal, 1, None)
370 .is_ok());
371 assert!(s
372 .submit("b", IODirection::Read, IOPriority::Normal, 1, None)
373 .is_ok());
374 let res = s.submit("c", IODirection::Read, IOPriority::Normal, 1, None);
375 assert!(res.is_err());
376 assert!(res.err().unwrap_or_default().contains("full"));
377 }
378
379 #[test]
382 fn test_next_request_priority_ordering() {
383 let mut s = default_scheduler();
384 let _ = s.submit("lo", IODirection::Read, IOPriority::Background, 1, None);
385 let _ = s.submit("hi", IODirection::Read, IOPriority::High, 1, None);
386 let _ = s.submit("mid", IODirection::Read, IOPriority::Normal, 1, None);
387
388 let r = s.next_request().expect("should have request");
389 assert_eq!(r.block_cid, "hi");
390 let r = s.next_request().expect("should have request");
391 assert_eq!(r.block_cid, "mid");
392 let r = s.next_request().expect("should have request");
393 assert_eq!(r.block_cid, "lo");
394 assert!(s.next_request().is_none());
395 }
396
397 #[test]
398 fn test_next_request_realtime_highest() {
399 let mut s = default_scheduler();
400 let _ = s.submit("h", IODirection::Read, IOPriority::High, 1, None);
401 let _ = s.submit("rt", IODirection::Read, IOPriority::Realtime, 1, None);
402
403 let r = s.next_request().expect("should have request");
404 assert_eq!(r.block_cid, "rt");
405 }
406
407 #[test]
410 fn test_fifo_within_same_priority() {
411 let mut s = default_scheduler();
412 let _ = s.submit("first", IODirection::Read, IOPriority::Normal, 1, None);
413 let _ = s.submit("second", IODirection::Read, IOPriority::Normal, 1, None);
414 let _ = s.submit("third", IODirection::Read, IOPriority::Normal, 1, None);
415
416 assert_eq!(s.next_request().expect("r").block_cid, "first");
417 assert_eq!(s.next_request().expect("r").block_cid, "second");
418 assert_eq!(s.next_request().expect("r").block_cid, "third");
419 }
420
421 #[test]
424 fn test_deadline_urgent_dispatched_first() {
425 let mut s = default_scheduler();
426 for _ in 0..5 {
428 s.tick();
429 }
430 let _ = s.submit("high", IODirection::Read, IOPriority::High, 1, None);
432 let _ = s.submit(
434 "urgent",
435 IODirection::Read,
436 IOPriority::Background,
437 1,
438 Some(3),
439 );
440
441 let r = s.next_request().expect("r");
442 assert_eq!(r.block_cid, "urgent");
443 }
444
445 #[test]
446 fn test_deadline_urgent_earliest_first() {
447 let mut s = default_scheduler();
448 for _ in 0..10 {
449 s.tick();
450 }
451 let _ = s.submit("dl5", IODirection::Read, IOPriority::Normal, 1, Some(5));
452 let _ = s.submit("dl3", IODirection::Read, IOPriority::Normal, 1, Some(3));
453 let _ = s.submit("dl7", IODirection::Read, IOPriority::Normal, 1, Some(7));
454
455 assert_eq!(s.next_request().expect("r").block_cid, "dl3");
456 assert_eq!(s.next_request().expect("r").block_cid, "dl5");
457 assert_eq!(s.next_request().expect("r").block_cid, "dl7");
458 }
459
460 #[test]
461 fn test_deadline_at_current_tick_is_urgent() {
462 let mut s = default_scheduler();
463 for _ in 0..5 {
464 s.tick();
465 }
466 let _ = s.submit(
468 "edge",
469 IODirection::Read,
470 IOPriority::Background,
471 1,
472 Some(5),
473 );
474 let _ = s.submit("nope", IODirection::Read, IOPriority::Realtime, 1, None);
475
476 assert_eq!(s.next_request().expect("r").block_cid, "edge");
477 }
478
479 #[test]
482 fn test_cancel_existing() {
483 let mut s = default_scheduler();
484 let id = s
485 .submit("c", IODirection::Read, IOPriority::Normal, 1, None)
486 .expect("ok");
487 assert_eq!(s.pending_count(), 1);
488 assert!(s.cancel(id));
489 assert_eq!(s.pending_count(), 0);
490 }
491
492 #[test]
493 fn test_cancel_nonexistent() {
494 let mut s = default_scheduler();
495 assert!(!s.cancel(999));
496 }
497
498 #[test]
499 fn test_cancel_double() {
500 let mut s = default_scheduler();
501 let id = s
502 .submit("c", IODirection::Read, IOPriority::Normal, 1, None)
503 .expect("ok");
504 assert!(s.cancel(id));
505 assert!(!s.cancel(id));
506 }
507
508 #[test]
511 fn test_complete_read_counting() {
512 let mut s = default_scheduler();
513 s.complete(0, IODirection::Read, 1024);
514 s.complete(1, IODirection::Read, 2048);
515 let st = s.stats();
516 assert_eq!(st.completed_reads, 2);
517 assert_eq!(st.completed_writes, 0);
518 assert_eq!(st.total_bytes_scheduled, 3072);
519 }
520
521 #[test]
522 fn test_complete_write_counting() {
523 let mut s = default_scheduler();
524 s.complete(0, IODirection::Write, 512);
525 let st = s.stats();
526 assert_eq!(st.completed_writes, 1);
527 assert_eq!(st.completed_reads, 0);
528 assert_eq!(st.total_bytes_scheduled, 512);
529 }
530
531 #[test]
532 fn test_complete_mixed() {
533 let mut s = default_scheduler();
534 s.complete(0, IODirection::Read, 100);
535 s.complete(1, IODirection::Write, 200);
536 s.complete(2, IODirection::Read, 300);
537 let st = s.stats();
538 assert_eq!(st.completed_reads, 2);
539 assert_eq!(st.completed_writes, 1);
540 assert_eq!(st.total_bytes_scheduled, 600);
541 }
542
543 #[test]
546 fn test_pending_reads_writes() {
547 let mut s = default_scheduler();
548 let _ = s.submit("r1", IODirection::Read, IOPriority::Normal, 1, None);
549 let _ = s.submit("r2", IODirection::Read, IOPriority::High, 1, None);
550 let _ = s.submit("w1", IODirection::Write, IOPriority::Normal, 1, None);
551
552 assert_eq!(s.pending_reads(), 2);
553 assert_eq!(s.pending_writes(), 1);
554 assert_eq!(s.pending_count(), 3);
555 }
556
557 #[test]
560 fn test_tick_advances_clock() {
561 let mut s = default_scheduler();
562 assert_eq!(s.current_tick(), 0);
563 s.tick();
564 assert_eq!(s.current_tick(), 1);
565 for _ in 0..10 {
566 s.tick();
567 }
568 assert_eq!(s.current_tick(), 11);
569 }
570
571 #[test]
574 fn test_expired_requests_none_initially() {
575 let mut s = default_scheduler();
576 let _ = s.submit("a", IODirection::Read, IOPriority::Normal, 1, Some(5));
577 assert!(s.expired_requests().is_empty());
578 }
579
580 #[test]
581 fn test_expired_requests_after_ticks() {
582 let mut s = default_scheduler();
583 let _ = s.submit("a", IODirection::Read, IOPriority::Normal, 1, Some(2));
584 let _ = s.submit("b", IODirection::Read, IOPriority::Normal, 1, Some(5));
585 let _ = s.submit("c", IODirection::Read, IOPriority::Normal, 1, None);
586
587 for _ in 0..3 {
589 s.tick();
590 }
591 let expired = s.expired_requests();
592 assert_eq!(expired.len(), 1);
593 assert_eq!(expired[0].block_cid, "a");
594 }
595
596 #[test]
597 fn test_expired_at_exact_tick_not_expired() {
598 let mut s = default_scheduler();
600 let _ = s.submit("a", IODirection::Read, IOPriority::Normal, 1, Some(3));
601 for _ in 0..3 {
602 s.tick();
603 }
604 assert!(s.expired_requests().is_empty());
606 }
607
608 #[test]
611 fn test_drain_expired() {
612 let mut s = default_scheduler();
613 let _ = s.submit("a", IODirection::Read, IOPriority::Normal, 1, Some(1));
614 let _ = s.submit("b", IODirection::Read, IOPriority::Normal, 1, Some(2));
615 let _ = s.submit("c", IODirection::Read, IOPriority::Normal, 1, Some(10));
616 let _ = s.submit("d", IODirection::Read, IOPriority::Normal, 1, None);
617
618 for _ in 0..5 {
619 s.tick();
620 }
621
622 let drained = s.drain_expired();
623 assert_eq!(drained.len(), 2);
624 assert_eq!(s.pending_count(), 2); }
626
627 #[test]
628 fn test_drain_expired_idempotent() {
629 let mut s = default_scheduler();
630 let _ = s.submit("a", IODirection::Read, IOPriority::Normal, 1, Some(1));
631 for _ in 0..3 {
632 s.tick();
633 }
634 let first = s.drain_expired();
635 assert_eq!(first.len(), 1);
636 let second = s.drain_expired();
637 assert!(second.is_empty());
638 }
639
640 #[test]
643 fn test_stats_accuracy() {
644 let mut s = default_scheduler();
645 let _ = s.submit("a", IODirection::Read, IOPriority::Normal, 100, Some(1));
646 let _ = s.submit("b", IODirection::Write, IOPriority::High, 200, None);
647 s.complete(10, IODirection::Read, 50);
648
649 for _ in 0..3 {
650 s.tick();
651 }
652
653 let st = s.stats();
654 assert_eq!(st.pending_count, 2);
655 assert_eq!(st.completed_reads, 1);
656 assert_eq!(st.completed_writes, 0);
657 assert_eq!(st.total_bytes_scheduled, 50);
658 assert_eq!(st.expired_count, 1); }
660
661 #[test]
664 fn test_mixed_read_write_ordering() {
665 let mut s = default_scheduler();
666 let _ = s.submit("w", IODirection::Write, IOPriority::High, 1, None);
667 let _ = s.submit("r", IODirection::Read, IOPriority::Normal, 1, None);
668
669 assert_eq!(s.next_request().expect("r").block_cid, "w");
671 assert_eq!(s.next_request().expect("r").block_cid, "r");
672 }
673
674 #[test]
677 fn test_next_request_empty() {
678 let mut s = default_scheduler();
679 assert!(s.next_request().is_none());
680 }
681
682 #[test]
685 fn test_submit_after_cancel() {
686 let mut s = small_scheduler(1);
687 let id = s
688 .submit("a", IODirection::Read, IOPriority::Normal, 1, None)
689 .expect("ok");
690 assert!(s
691 .submit("b", IODirection::Read, IOPriority::Normal, 1, None)
692 .is_err());
693 s.cancel(id);
694 assert!(s
695 .submit("c", IODirection::Read, IOPriority::Normal, 1, None)
696 .is_ok());
697 }
698
699 #[test]
702 fn test_large_batch_submit_and_drain() {
703 let mut s = default_scheduler();
704 for i in 0..100 {
705 let _ = s.submit(
706 &format!("blk{i}"),
707 IODirection::Read,
708 IOPriority::Normal,
709 64,
710 Some(50),
711 );
712 }
713 assert_eq!(s.pending_count(), 100);
714
715 for _ in 0..60 {
716 s.tick();
717 }
718 let drained = s.drain_expired();
719 assert_eq!(drained.len(), 100);
720 assert_eq!(s.pending_count(), 0);
721 }
722
723 #[test]
726 fn test_config_accessor() {
727 let s = default_scheduler();
728 assert_eq!(s.config().max_pending, 1000);
729 assert!((s.config().read_weight - 0.7).abs() < 1e-9);
730 assert!((s.config().write_weight - 0.3).abs() < 1e-9);
731 }
732
733 #[test]
736 fn test_default_config() {
737 let cfg = SchedulerConfig::default();
738 assert_eq!(cfg.max_pending, 1000);
739 assert!((cfg.read_weight - 0.7).abs() < 1e-9);
740 assert!((cfg.write_weight - 0.3).abs() < 1e-9);
741 }
742
743 #[test]
746 fn test_complete_removes_from_pending() {
747 let mut s = default_scheduler();
748 let id = s
749 .submit("x", IODirection::Read, IOPriority::Normal, 256, None)
750 .expect("ok");
751 assert_eq!(s.pending_count(), 1);
752 s.complete(id, IODirection::Read, 256);
753 assert_eq!(s.pending_count(), 0);
754 assert_eq!(s.stats().completed_reads, 1);
755 }
756
757 #[test]
760 fn test_deadline_urgent_fifo_tiebreak() {
761 let mut s = default_scheduler();
762 let _ = s.submit("first", IODirection::Read, IOPriority::Normal, 1, Some(2));
764 let _ = s.submit("second", IODirection::Read, IOPriority::Normal, 1, Some(2));
765
766 for _ in 0..3 {
767 s.tick();
768 }
769 assert_eq!(s.next_request().expect("r").block_cid, "first");
770 assert_eq!(s.next_request().expect("r").block_cid, "second");
771 }
772
773 #[test]
776 fn test_non_urgent_deadline_does_not_promote() {
777 let mut s = default_scheduler();
778 let _ = s.submit(
780 "dl_future",
781 IODirection::Read,
782 IOPriority::Background,
783 1,
784 Some(100),
785 );
786 let _ = s.submit("high", IODirection::Read, IOPriority::High, 1, None);
787
788 assert_eq!(s.next_request().expect("r").block_cid, "high");
790 }
791}