#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum IOPriority {
Background = 0,
Normal = 1,
High = 2,
Realtime = 3,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum IODirection {
Read,
Write,
}
#[derive(Debug, Clone)]
pub struct IORequest {
pub id: u64,
pub block_cid: String,
pub direction: IODirection,
pub priority: IOPriority,
pub size_bytes: u64,
pub deadline_tick: Option<u64>,
pub enqueued_tick: u64,
}
#[derive(Debug, Clone)]
pub struct SchedulerConfig {
pub max_pending: usize,
pub read_weight: f64,
pub write_weight: f64,
}
impl Default for SchedulerConfig {
fn default() -> Self {
Self {
max_pending: 1000,
read_weight: 0.7,
write_weight: 0.3,
}
}
}
#[derive(Debug, Clone)]
pub struct IOSchedulerStats {
pub pending_count: usize,
pub completed_reads: u64,
pub completed_writes: u64,
pub total_bytes_scheduled: u64,
pub expired_count: usize,
}
pub struct StorageIOScheduler {
config: SchedulerConfig,
pending: Vec<IORequest>,
next_id: u64,
current_tick: u64,
completed_reads: u64,
completed_writes: u64,
total_bytes_scheduled: u64,
}
impl StorageIOScheduler {
pub fn new(config: SchedulerConfig) -> Self {
Self {
config,
pending: Vec::new(),
next_id: 0,
current_tick: 0,
completed_reads: 0,
completed_writes: 0,
total_bytes_scheduled: 0,
}
}
pub fn submit(
&mut self,
block_cid: &str,
direction: IODirection,
priority: IOPriority,
size_bytes: u64,
deadline_tick: Option<u64>,
) -> Result<u64, String> {
if self.pending.len() >= self.config.max_pending {
return Err(format!(
"pending queue full ({} / {})",
self.pending.len(),
self.config.max_pending,
));
}
let id = self.next_id;
self.next_id = self.next_id.saturating_add(1);
self.pending.push(IORequest {
id,
block_cid: block_cid.to_string(),
direction,
priority,
size_bytes,
deadline_tick,
enqueued_tick: self.current_tick,
});
Ok(id)
}
pub fn next_request(&mut self) -> Option<IORequest> {
if self.pending.is_empty() {
return None;
}
let tick = self.current_tick;
let mut best_idx: usize = 0;
let mut best_urgent = false;
let mut best_deadline: u64 = u64::MAX;
let mut best_priority = IOPriority::Background;
let mut best_enqueued: u64 = u64::MAX;
for (i, req) in self.pending.iter().enumerate() {
let urgent = req.deadline_tick.map(|d| d <= tick).unwrap_or(false);
let better = if urgent && !best_urgent {
true
} else if urgent && best_urgent {
let dl = req.deadline_tick.unwrap_or(u64::MAX);
dl < best_deadline || (dl == best_deadline && req.enqueued_tick < best_enqueued)
} else if !urgent && best_urgent {
false
} else {
req.priority > best_priority
|| (req.priority == best_priority && req.enqueued_tick < best_enqueued)
};
if better {
best_idx = i;
best_urgent = urgent;
best_deadline = req.deadline_tick.unwrap_or(u64::MAX);
best_priority = req.priority;
best_enqueued = req.enqueued_tick;
}
}
Some(self.pending.remove(best_idx))
}
pub fn complete(&mut self, request_id: u64, direction: IODirection, size_bytes: u64) {
match direction {
IODirection::Read => self.completed_reads = self.completed_reads.saturating_add(1),
IODirection::Write => self.completed_writes = self.completed_writes.saturating_add(1),
}
self.total_bytes_scheduled = self.total_bytes_scheduled.saturating_add(size_bytes);
self.pending.retain(|r| r.id != request_id);
}
pub fn cancel(&mut self, request_id: u64) -> bool {
let before = self.pending.len();
self.pending.retain(|r| r.id != request_id);
self.pending.len() < before
}
pub fn pending_count(&self) -> usize {
self.pending.len()
}
pub fn pending_reads(&self) -> usize {
self.pending
.iter()
.filter(|r| r.direction == IODirection::Read)
.count()
}
pub fn pending_writes(&self) -> usize {
self.pending
.iter()
.filter(|r| r.direction == IODirection::Write)
.count()
}
pub fn tick(&mut self) {
self.current_tick = self.current_tick.saturating_add(1);
}
pub fn current_tick(&self) -> u64 {
self.current_tick
}
pub fn expired_requests(&self) -> Vec<&IORequest> {
let tick = self.current_tick;
self.pending
.iter()
.filter(|r| r.deadline_tick.map(|d| d < tick).unwrap_or(false))
.collect()
}
pub fn drain_expired(&mut self) -> Vec<IORequest> {
let tick = self.current_tick;
let mut expired = Vec::new();
let mut kept = Vec::new();
for req in self.pending.drain(..) {
if req.deadline_tick.map(|d| d < tick).unwrap_or(false) {
expired.push(req);
} else {
kept.push(req);
}
}
self.pending = kept;
expired
}
pub fn stats(&self) -> IOSchedulerStats {
IOSchedulerStats {
pending_count: self.pending.len(),
completed_reads: self.completed_reads,
completed_writes: self.completed_writes,
total_bytes_scheduled: self.total_bytes_scheduled,
expired_count: self.expired_requests().len(),
}
}
pub fn config(&self) -> &SchedulerConfig {
&self.config
}
}
#[cfg(test)]
mod tests {
use super::*;
fn default_scheduler() -> StorageIOScheduler {
StorageIOScheduler::new(SchedulerConfig::default())
}
fn small_scheduler(max: usize) -> StorageIOScheduler {
StorageIOScheduler::new(SchedulerConfig {
max_pending: max,
..SchedulerConfig::default()
})
}
#[test]
fn test_submit_returns_unique_ids() {
let mut s = default_scheduler();
let id0 = s.submit("cid0", IODirection::Read, IOPriority::Normal, 100, None);
let id1 = s.submit("cid1", IODirection::Write, IOPriority::Normal, 200, None);
assert!(id0.is_ok());
assert!(id1.is_ok());
assert_ne!(id0.ok(), id1.ok());
}
#[test]
fn test_submit_increments_pending() {
let mut s = default_scheduler();
assert_eq!(s.pending_count(), 0);
let _ = s.submit("c", IODirection::Read, IOPriority::Normal, 10, None);
assert_eq!(s.pending_count(), 1);
let _ = s.submit("c", IODirection::Write, IOPriority::High, 20, None);
assert_eq!(s.pending_count(), 2);
}
#[test]
fn test_submit_max_pending_enforced() {
let mut s = small_scheduler(2);
assert!(s
.submit("a", IODirection::Read, IOPriority::Normal, 1, None)
.is_ok());
assert!(s
.submit("b", IODirection::Read, IOPriority::Normal, 1, None)
.is_ok());
let res = s.submit("c", IODirection::Read, IOPriority::Normal, 1, None);
assert!(res.is_err());
assert!(res.err().unwrap_or_default().contains("full"));
}
#[test]
fn test_next_request_priority_ordering() {
let mut s = default_scheduler();
let _ = s.submit("lo", IODirection::Read, IOPriority::Background, 1, None);
let _ = s.submit("hi", IODirection::Read, IOPriority::High, 1, None);
let _ = s.submit("mid", IODirection::Read, IOPriority::Normal, 1, None);
let r = s.next_request().expect("should have request");
assert_eq!(r.block_cid, "hi");
let r = s.next_request().expect("should have request");
assert_eq!(r.block_cid, "mid");
let r = s.next_request().expect("should have request");
assert_eq!(r.block_cid, "lo");
assert!(s.next_request().is_none());
}
#[test]
fn test_next_request_realtime_highest() {
let mut s = default_scheduler();
let _ = s.submit("h", IODirection::Read, IOPriority::High, 1, None);
let _ = s.submit("rt", IODirection::Read, IOPriority::Realtime, 1, None);
let r = s.next_request().expect("should have request");
assert_eq!(r.block_cid, "rt");
}
#[test]
fn test_fifo_within_same_priority() {
let mut s = default_scheduler();
let _ = s.submit("first", IODirection::Read, IOPriority::Normal, 1, None);
let _ = s.submit("second", IODirection::Read, IOPriority::Normal, 1, None);
let _ = s.submit("third", IODirection::Read, IOPriority::Normal, 1, None);
assert_eq!(s.next_request().expect("r").block_cid, "first");
assert_eq!(s.next_request().expect("r").block_cid, "second");
assert_eq!(s.next_request().expect("r").block_cid, "third");
}
#[test]
fn test_deadline_urgent_dispatched_first() {
let mut s = default_scheduler();
for _ in 0..5 {
s.tick();
}
let _ = s.submit("high", IODirection::Read, IOPriority::High, 1, None);
let _ = s.submit(
"urgent",
IODirection::Read,
IOPriority::Background,
1,
Some(3),
);
let r = s.next_request().expect("r");
assert_eq!(r.block_cid, "urgent");
}
#[test]
fn test_deadline_urgent_earliest_first() {
let mut s = default_scheduler();
for _ in 0..10 {
s.tick();
}
let _ = s.submit("dl5", IODirection::Read, IOPriority::Normal, 1, Some(5));
let _ = s.submit("dl3", IODirection::Read, IOPriority::Normal, 1, Some(3));
let _ = s.submit("dl7", IODirection::Read, IOPriority::Normal, 1, Some(7));
assert_eq!(s.next_request().expect("r").block_cid, "dl3");
assert_eq!(s.next_request().expect("r").block_cid, "dl5");
assert_eq!(s.next_request().expect("r").block_cid, "dl7");
}
#[test]
fn test_deadline_at_current_tick_is_urgent() {
let mut s = default_scheduler();
for _ in 0..5 {
s.tick();
}
let _ = s.submit(
"edge",
IODirection::Read,
IOPriority::Background,
1,
Some(5),
);
let _ = s.submit("nope", IODirection::Read, IOPriority::Realtime, 1, None);
assert_eq!(s.next_request().expect("r").block_cid, "edge");
}
#[test]
fn test_cancel_existing() {
let mut s = default_scheduler();
let id = s
.submit("c", IODirection::Read, IOPriority::Normal, 1, None)
.expect("ok");
assert_eq!(s.pending_count(), 1);
assert!(s.cancel(id));
assert_eq!(s.pending_count(), 0);
}
#[test]
fn test_cancel_nonexistent() {
let mut s = default_scheduler();
assert!(!s.cancel(999));
}
#[test]
fn test_cancel_double() {
let mut s = default_scheduler();
let id = s
.submit("c", IODirection::Read, IOPriority::Normal, 1, None)
.expect("ok");
assert!(s.cancel(id));
assert!(!s.cancel(id));
}
#[test]
fn test_complete_read_counting() {
let mut s = default_scheduler();
s.complete(0, IODirection::Read, 1024);
s.complete(1, IODirection::Read, 2048);
let st = s.stats();
assert_eq!(st.completed_reads, 2);
assert_eq!(st.completed_writes, 0);
assert_eq!(st.total_bytes_scheduled, 3072);
}
#[test]
fn test_complete_write_counting() {
let mut s = default_scheduler();
s.complete(0, IODirection::Write, 512);
let st = s.stats();
assert_eq!(st.completed_writes, 1);
assert_eq!(st.completed_reads, 0);
assert_eq!(st.total_bytes_scheduled, 512);
}
#[test]
fn test_complete_mixed() {
let mut s = default_scheduler();
s.complete(0, IODirection::Read, 100);
s.complete(1, IODirection::Write, 200);
s.complete(2, IODirection::Read, 300);
let st = s.stats();
assert_eq!(st.completed_reads, 2);
assert_eq!(st.completed_writes, 1);
assert_eq!(st.total_bytes_scheduled, 600);
}
#[test]
fn test_pending_reads_writes() {
let mut s = default_scheduler();
let _ = s.submit("r1", IODirection::Read, IOPriority::Normal, 1, None);
let _ = s.submit("r2", IODirection::Read, IOPriority::High, 1, None);
let _ = s.submit("w1", IODirection::Write, IOPriority::Normal, 1, None);
assert_eq!(s.pending_reads(), 2);
assert_eq!(s.pending_writes(), 1);
assert_eq!(s.pending_count(), 3);
}
#[test]
fn test_tick_advances_clock() {
let mut s = default_scheduler();
assert_eq!(s.current_tick(), 0);
s.tick();
assert_eq!(s.current_tick(), 1);
for _ in 0..10 {
s.tick();
}
assert_eq!(s.current_tick(), 11);
}
#[test]
fn test_expired_requests_none_initially() {
let mut s = default_scheduler();
let _ = s.submit("a", IODirection::Read, IOPriority::Normal, 1, Some(5));
assert!(s.expired_requests().is_empty());
}
#[test]
fn test_expired_requests_after_ticks() {
let mut s = default_scheduler();
let _ = s.submit("a", IODirection::Read, IOPriority::Normal, 1, Some(2));
let _ = s.submit("b", IODirection::Read, IOPriority::Normal, 1, Some(5));
let _ = s.submit("c", IODirection::Read, IOPriority::Normal, 1, None);
for _ in 0..3 {
s.tick();
}
let expired = s.expired_requests();
assert_eq!(expired.len(), 1);
assert_eq!(expired[0].block_cid, "a");
}
#[test]
fn test_expired_at_exact_tick_not_expired() {
let mut s = default_scheduler();
let _ = s.submit("a", IODirection::Read, IOPriority::Normal, 1, Some(3));
for _ in 0..3 {
s.tick();
}
assert!(s.expired_requests().is_empty());
}
#[test]
fn test_drain_expired() {
let mut s = default_scheduler();
let _ = s.submit("a", IODirection::Read, IOPriority::Normal, 1, Some(1));
let _ = s.submit("b", IODirection::Read, IOPriority::Normal, 1, Some(2));
let _ = s.submit("c", IODirection::Read, IOPriority::Normal, 1, Some(10));
let _ = s.submit("d", IODirection::Read, IOPriority::Normal, 1, None);
for _ in 0..5 {
s.tick();
}
let drained = s.drain_expired();
assert_eq!(drained.len(), 2);
assert_eq!(s.pending_count(), 2); }
#[test]
fn test_drain_expired_idempotent() {
let mut s = default_scheduler();
let _ = s.submit("a", IODirection::Read, IOPriority::Normal, 1, Some(1));
for _ in 0..3 {
s.tick();
}
let first = s.drain_expired();
assert_eq!(first.len(), 1);
let second = s.drain_expired();
assert!(second.is_empty());
}
#[test]
fn test_stats_accuracy() {
let mut s = default_scheduler();
let _ = s.submit("a", IODirection::Read, IOPriority::Normal, 100, Some(1));
let _ = s.submit("b", IODirection::Write, IOPriority::High, 200, None);
s.complete(10, IODirection::Read, 50);
for _ in 0..3 {
s.tick();
}
let st = s.stats();
assert_eq!(st.pending_count, 2);
assert_eq!(st.completed_reads, 1);
assert_eq!(st.completed_writes, 0);
assert_eq!(st.total_bytes_scheduled, 50);
assert_eq!(st.expired_count, 1); }
#[test]
fn test_mixed_read_write_ordering() {
let mut s = default_scheduler();
let _ = s.submit("w", IODirection::Write, IOPriority::High, 1, None);
let _ = s.submit("r", IODirection::Read, IOPriority::Normal, 1, None);
assert_eq!(s.next_request().expect("r").block_cid, "w");
assert_eq!(s.next_request().expect("r").block_cid, "r");
}
#[test]
fn test_next_request_empty() {
let mut s = default_scheduler();
assert!(s.next_request().is_none());
}
#[test]
fn test_submit_after_cancel() {
let mut s = small_scheduler(1);
let id = s
.submit("a", IODirection::Read, IOPriority::Normal, 1, None)
.expect("ok");
assert!(s
.submit("b", IODirection::Read, IOPriority::Normal, 1, None)
.is_err());
s.cancel(id);
assert!(s
.submit("c", IODirection::Read, IOPriority::Normal, 1, None)
.is_ok());
}
#[test]
fn test_large_batch_submit_and_drain() {
let mut s = default_scheduler();
for i in 0..100 {
let _ = s.submit(
&format!("blk{i}"),
IODirection::Read,
IOPriority::Normal,
64,
Some(50),
);
}
assert_eq!(s.pending_count(), 100);
for _ in 0..60 {
s.tick();
}
let drained = s.drain_expired();
assert_eq!(drained.len(), 100);
assert_eq!(s.pending_count(), 0);
}
#[test]
fn test_config_accessor() {
let s = default_scheduler();
assert_eq!(s.config().max_pending, 1000);
assert!((s.config().read_weight - 0.7).abs() < 1e-9);
assert!((s.config().write_weight - 0.3).abs() < 1e-9);
}
#[test]
fn test_default_config() {
let cfg = SchedulerConfig::default();
assert_eq!(cfg.max_pending, 1000);
assert!((cfg.read_weight - 0.7).abs() < 1e-9);
assert!((cfg.write_weight - 0.3).abs() < 1e-9);
}
#[test]
fn test_complete_removes_from_pending() {
let mut s = default_scheduler();
let id = s
.submit("x", IODirection::Read, IOPriority::Normal, 256, None)
.expect("ok");
assert_eq!(s.pending_count(), 1);
s.complete(id, IODirection::Read, 256);
assert_eq!(s.pending_count(), 0);
assert_eq!(s.stats().completed_reads, 1);
}
#[test]
fn test_deadline_urgent_fifo_tiebreak() {
let mut s = default_scheduler();
let _ = s.submit("first", IODirection::Read, IOPriority::Normal, 1, Some(2));
let _ = s.submit("second", IODirection::Read, IOPriority::Normal, 1, Some(2));
for _ in 0..3 {
s.tick();
}
assert_eq!(s.next_request().expect("r").block_cid, "first");
assert_eq!(s.next_request().expect("r").block_cid, "second");
}
#[test]
fn test_non_urgent_deadline_does_not_promote() {
let mut s = default_scheduler();
let _ = s.submit(
"dl_future",
IODirection::Read,
IOPriority::Background,
1,
Some(100),
);
let _ = s.submit("high", IODirection::Read, IOPriority::High, 1, None);
assert_eq!(s.next_request().expect("r").block_cid, "high");
}
}