use std::collections::HashMap;
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum StorageTier {
Nvme = 0,
Ssd = 1,
Hdd = 2,
Archive = 3,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MigrationDirection {
Promote,
Demote,
}
#[derive(Clone, Debug)]
pub struct MigrationTask {
pub task_id: u64,
pub block_id: u64,
pub cid: String,
pub from_tier: StorageTier,
pub to_tier: StorageTier,
pub direction: MigrationDirection,
pub size_bytes: u64,
pub estimated_cost_ms: u64,
pub depends_on: Vec<u64>,
}
impl MigrationTask {
pub fn is_promotion(&self) -> bool {
self.direction == MigrationDirection::Promote
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MigrationStatus {
Pending,
InProgress,
Completed,
Failed,
RolledBack,
}
#[derive(Clone, Debug)]
pub struct MigrationRecord {
pub task: MigrationTask,
pub status: MigrationStatus,
pub started_at_tick: Option<u64>,
pub completed_at_tick: Option<u64>,
}
#[derive(Clone, Debug, Default)]
pub struct PlannerStats {
pub total_tasks: usize,
pub pending: usize,
pub in_progress: usize,
pub completed: usize,
pub failed: usize,
pub total_bytes_migrated: u64,
pub promotions: u64,
pub demotions: u64,
}
pub struct StorageMigrationPlanner {
pub records: HashMap<u64, MigrationRecord>,
pub next_task_id: u64,
}
impl StorageMigrationPlanner {
pub fn new() -> Self {
Self {
records: HashMap::new(),
next_task_id: 0,
}
}
pub fn plan_migration(
&mut self,
block_id: u64,
cid: String,
from: StorageTier,
to: StorageTier,
size_bytes: u64,
depends_on: Vec<u64>,
) -> u64 {
let task_id = self.next_task_id;
self.next_task_id += 1;
let direction = if from > to {
MigrationDirection::Promote
} else {
MigrationDirection::Demote
};
let estimated_cost_ms = size_bytes / 1_000_000 + 10;
let task = MigrationTask {
task_id,
block_id,
cid,
from_tier: from,
to_tier: to,
direction,
size_bytes,
estimated_cost_ms,
depends_on,
};
let record = MigrationRecord {
task,
status: MigrationStatus::Pending,
started_at_tick: None,
completed_at_tick: None,
};
self.records.insert(task_id, record);
task_id
}
pub fn start_task(&mut self, task_id: u64, current_tick: u64) -> bool {
let depends_on: Vec<u64> = match self.records.get(&task_id) {
Some(r) if r.status == MigrationStatus::Pending => r.task.depends_on.clone(),
_ => return false,
};
for dep_id in &depends_on {
match self.records.get(dep_id) {
Some(dep) if dep.status == MigrationStatus::Completed => {}
_ => return false,
}
}
if let Some(record) = self.records.get_mut(&task_id) {
record.status = MigrationStatus::InProgress;
record.started_at_tick = Some(current_tick);
true
} else {
false
}
}
pub fn complete_task(&mut self, task_id: u64, current_tick: u64) -> bool {
match self.records.get_mut(&task_id) {
Some(r) if r.status == MigrationStatus::InProgress => {
r.status = MigrationStatus::Completed;
r.completed_at_tick = Some(current_tick);
true
}
_ => false,
}
}
pub fn fail_task(&mut self, task_id: u64) -> bool {
match self.records.get_mut(&task_id) {
Some(r)
if r.status == MigrationStatus::Pending
|| r.status == MigrationStatus::InProgress =>
{
r.status = MigrationStatus::Failed;
true
}
_ => false,
}
}
pub fn rollback_task(&mut self, task_id: u64) -> bool {
match self.records.get_mut(&task_id) {
Some(r) if r.status == MigrationStatus::Failed => {
r.status = MigrationStatus::RolledBack;
true
}
_ => false,
}
}
pub fn ready_tasks(&self) -> Vec<u64> {
let mut ready: Vec<u64> = self
.records
.values()
.filter(|r| {
if r.status != MigrationStatus::Pending {
return false;
}
r.task.depends_on.iter().all(|dep_id| {
self.records
.get(dep_id)
.is_some_and(|d| d.status == MigrationStatus::Completed)
})
})
.map(|r| r.task.task_id)
.collect();
ready.sort_unstable();
ready
}
pub fn get_record(&self, task_id: u64) -> Option<&MigrationRecord> {
self.records.get(&task_id)
}
pub fn stats(&self) -> PlannerStats {
let mut stats = PlannerStats {
total_tasks: self.records.len(),
..Default::default()
};
for record in self.records.values() {
match record.status {
MigrationStatus::Pending => stats.pending += 1,
MigrationStatus::InProgress => stats.in_progress += 1,
MigrationStatus::Completed => {
stats.completed += 1;
stats.total_bytes_migrated += record.task.size_bytes;
match record.task.direction {
MigrationDirection::Promote => stats.promotions += 1,
MigrationDirection::Demote => stats.demotions += 1,
}
}
MigrationStatus::Failed | MigrationStatus::RolledBack => stats.failed += 1,
}
}
stats
}
}
impl Default for StorageMigrationPlanner {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_planner() -> StorageMigrationPlanner {
StorageMigrationPlanner::new()
}
fn add_simple(
p: &mut StorageMigrationPlanner,
from: StorageTier,
to: StorageTier,
size: u64,
) -> u64 {
p.plan_migration(1, "bafytest".into(), from, to, size, vec![])
}
#[test]
fn plan_migration_returns_sequential_ids() {
let mut p = make_planner();
let id0 = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
let id1 = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
assert_eq!(id0, 0);
assert_eq!(id1, 1);
}
#[test]
fn plan_migration_demote_direction() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
let record = p.get_record(id).expect("record must exist");
assert_eq!(record.task.direction, MigrationDirection::Demote);
assert!(!record.task.is_promotion());
}
#[test]
fn plan_migration_promote_direction() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Archive, StorageTier::Ssd, 0);
let record = p.get_record(id).expect("record must exist");
assert_eq!(record.task.direction, MigrationDirection::Promote);
assert!(record.task.is_promotion());
}
#[test]
fn plan_migration_same_tier_is_demote() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Ssd, StorageTier::Ssd, 0);
let record = p.get_record(id).expect("record must exist");
assert_eq!(record.task.direction, MigrationDirection::Demote);
}
#[test]
fn plan_migration_cost_estimate_small() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Nvme, StorageTier::Archive, 500_000);
let record = p.get_record(id).unwrap();
assert_eq!(record.task.estimated_cost_ms, 10);
}
#[test]
fn plan_migration_cost_estimate_large() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Ssd, StorageTier::Hdd, 5_000_000);
let record = p.get_record(id).unwrap();
assert_eq!(record.task.estimated_cost_ms, 15);
}
#[test]
fn plan_migration_status_is_pending() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
let record = p.get_record(id).unwrap();
assert_eq!(record.status, MigrationStatus::Pending);
assert!(record.started_at_tick.is_none());
assert!(record.completed_at_tick.is_none());
}
#[test]
fn plan_migration_stores_block_fields() {
let mut p = make_planner();
let id = p.plan_migration(
42,
"bafycid".into(),
StorageTier::Hdd,
StorageTier::Nvme,
1_000,
vec![],
);
let record = p.get_record(id).unwrap();
assert_eq!(record.task.block_id, 42);
assert_eq!(record.task.cid, "bafycid");
assert_eq!(record.task.from_tier, StorageTier::Hdd);
assert_eq!(record.task.to_tier, StorageTier::Nvme);
assert_eq!(record.task.size_bytes, 1_000);
}
#[test]
fn start_task_sets_in_progress() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Ssd, StorageTier::Hdd, 0);
assert!(p.start_task(id, 5));
let record = p.get_record(id).unwrap();
assert_eq!(record.status, MigrationStatus::InProgress);
assert_eq!(record.started_at_tick, Some(5));
}
#[test]
fn start_task_unknown_id_returns_false() {
let mut p = make_planner();
assert!(!p.start_task(999, 0));
}
#[test]
fn start_task_already_in_progress_returns_false() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Ssd, StorageTier::Hdd, 0);
assert!(p.start_task(id, 1));
assert!(!p.start_task(id, 2));
}
#[test]
fn start_task_blocked_by_pending_dependency() {
let mut p = make_planner();
let dep = add_simple(&mut p, StorageTier::Nvme, StorageTier::Ssd, 0);
let child = p.plan_migration(
2,
"baf".into(),
StorageTier::Ssd,
StorageTier::Hdd,
0,
vec![dep],
);
assert!(!p.start_task(child, 0));
}
#[test]
fn start_task_blocked_by_in_progress_dependency() {
let mut p = make_planner();
let dep = add_simple(&mut p, StorageTier::Nvme, StorageTier::Ssd, 0);
let child = p.plan_migration(
2,
"baf".into(),
StorageTier::Ssd,
StorageTier::Hdd,
0,
vec![dep],
);
p.start_task(dep, 0);
assert!(!p.start_task(child, 1));
}
#[test]
fn start_task_unblocked_when_dependency_completes() {
let mut p = make_planner();
let dep = add_simple(&mut p, StorageTier::Nvme, StorageTier::Ssd, 0);
let child = p.plan_migration(
2,
"baf".into(),
StorageTier::Ssd,
StorageTier::Hdd,
0,
vec![dep],
);
p.start_task(dep, 0);
p.complete_task(dep, 1);
assert!(p.start_task(child, 2));
}
#[test]
fn complete_task_sets_completed() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Ssd, StorageTier::Archive, 0);
p.start_task(id, 0);
assert!(p.complete_task(id, 10));
let record = p.get_record(id).unwrap();
assert_eq!(record.status, MigrationStatus::Completed);
assert_eq!(record.completed_at_tick, Some(10));
}
#[test]
fn complete_task_not_in_progress_returns_false() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Ssd, StorageTier::Archive, 0);
assert!(!p.complete_task(id, 1));
}
#[test]
fn complete_task_unknown_id_returns_false() {
let mut p = make_planner();
assert!(!p.complete_task(999, 0));
}
#[test]
fn fail_task_from_pending() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
assert!(p.fail_task(id));
assert_eq!(p.get_record(id).unwrap().status, MigrationStatus::Failed);
}
#[test]
fn fail_task_from_in_progress() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
p.start_task(id, 0);
assert!(p.fail_task(id));
assert_eq!(p.get_record(id).unwrap().status, MigrationStatus::Failed);
}
#[test]
fn fail_task_from_completed_returns_false() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
p.start_task(id, 0);
p.complete_task(id, 1);
assert!(!p.fail_task(id));
}
#[test]
fn fail_task_unknown_id_returns_false() {
let mut p = make_planner();
assert!(!p.fail_task(999));
}
#[test]
fn rollback_task_from_failed_succeeds() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
p.fail_task(id);
assert!(p.rollback_task(id));
assert_eq!(
p.get_record(id).unwrap().status,
MigrationStatus::RolledBack
);
}
#[test]
fn rollback_task_from_pending_returns_false() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
assert!(!p.rollback_task(id));
}
#[test]
fn rollback_task_from_completed_returns_false() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
p.start_task(id, 0);
p.complete_task(id, 1);
assert!(!p.rollback_task(id));
}
#[test]
fn rollback_task_idempotency_returns_false() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
p.fail_task(id);
p.rollback_task(id);
assert!(!p.rollback_task(id));
}
#[test]
fn ready_tasks_no_deps_all_pending() {
let mut p = make_planner();
let id0 = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
let id1 = add_simple(&mut p, StorageTier::Ssd, StorageTier::Archive, 0);
let ready = p.ready_tasks();
assert_eq!(ready, vec![id0, id1]);
}
#[test]
fn ready_tasks_sorted_ascending() {
let mut p = make_planner();
for _ in 0..5 {
add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
}
let ready = p.ready_tasks();
let mut sorted = ready.clone();
sorted.sort_unstable();
assert_eq!(ready, sorted);
}
#[test]
fn ready_tasks_excludes_in_progress() {
let mut p = make_planner();
let id0 = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
let id1 = add_simple(&mut p, StorageTier::Ssd, StorageTier::Archive, 0);
p.start_task(id0, 0);
let ready = p.ready_tasks();
assert_eq!(ready, vec![id1]);
}
#[test]
fn ready_tasks_excludes_completed() {
let mut p = make_planner();
let id0 = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
p.start_task(id0, 0);
p.complete_task(id0, 1);
assert!(p.ready_tasks().is_empty());
}
#[test]
fn ready_tasks_with_dependency_chain() {
let mut p = make_planner();
let a = add_simple(&mut p, StorageTier::Nvme, StorageTier::Ssd, 0);
let b = p.plan_migration(
2,
"b".into(),
StorageTier::Ssd,
StorageTier::Hdd,
0,
vec![a],
);
let c = p.plan_migration(
3,
"c".into(),
StorageTier::Hdd,
StorageTier::Archive,
0,
vec![b],
);
assert_eq!(p.ready_tasks(), vec![a]);
p.start_task(a, 0);
p.complete_task(a, 1);
assert_eq!(p.ready_tasks(), vec![b]);
p.start_task(b, 2);
p.complete_task(b, 3);
assert_eq!(p.ready_tasks(), vec![c]);
}
#[test]
fn stats_counts_by_status() {
let mut p = make_planner();
let a = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 1_000);
let b = add_simple(&mut p, StorageTier::Ssd, StorageTier::Archive, 2_000);
let c = add_simple(&mut p, StorageTier::Hdd, StorageTier::Nvme, 3_000);
let d = add_simple(&mut p, StorageTier::Archive, StorageTier::Ssd, 500);
p.start_task(a, 0);
p.complete_task(a, 1);
p.start_task(b, 2);
p.fail_task(b);
let s = p.stats();
assert_eq!(s.total_tasks, 4);
assert_eq!(s.completed, 1);
assert_eq!(s.failed, 1);
assert_eq!(s.pending, 2);
assert_eq!(s.in_progress, 0);
let _ = c;
let _ = d;
}
#[test]
fn stats_total_bytes_migrated_accumulates() {
let mut p = make_planner();
let a = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 1_000_000);
let b = add_simple(&mut p, StorageTier::Ssd, StorageTier::Archive, 2_000_000);
p.start_task(a, 0);
p.complete_task(a, 1);
p.start_task(b, 2);
p.complete_task(b, 3);
assert_eq!(p.stats().total_bytes_migrated, 3_000_000);
}
#[test]
fn stats_bytes_not_counted_for_failed() {
let mut p = make_planner();
let a = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 5_000_000);
p.start_task(a, 0);
p.fail_task(a);
assert_eq!(p.stats().total_bytes_migrated, 0);
}
#[test]
fn stats_promotions_and_demotions() {
let mut p = make_planner();
let d0 = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
let d1 = add_simple(&mut p, StorageTier::Ssd, StorageTier::Archive, 0);
let p0 = add_simple(&mut p, StorageTier::Hdd, StorageTier::Ssd, 0);
for id in [d0, d1, p0] {
p.start_task(id, 0);
p.complete_task(id, 1);
}
let s = p.stats();
assert_eq!(s.demotions, 2);
assert_eq!(s.promotions, 1);
}
#[test]
fn stats_rolledback_counted_as_failed() {
let mut p = make_planner();
let id = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
p.fail_task(id);
p.rollback_task(id);
let s = p.stats();
assert_eq!(s.failed, 1);
assert_eq!(s.completed, 0);
}
#[test]
fn stats_empty_planner() {
let p = make_planner();
let s = p.stats();
assert_eq!(s.total_tasks, 0);
assert_eq!(s.pending, 0);
assert_eq!(s.total_bytes_migrated, 0);
}
#[test]
fn get_record_missing_returns_none() {
let p = make_planner();
assert!(p.get_record(0).is_none());
}
#[test]
fn multiple_deps_all_must_complete() {
let mut p = make_planner();
let a = add_simple(&mut p, StorageTier::Nvme, StorageTier::Ssd, 0);
let b = add_simple(&mut p, StorageTier::Nvme, StorageTier::Hdd, 0);
let c = p.plan_migration(
10,
"c".into(),
StorageTier::Ssd,
StorageTier::Archive,
0,
vec![a, b],
);
assert!(!p.ready_tasks().contains(&c));
p.start_task(a, 0);
p.complete_task(a, 1);
assert!(!p.ready_tasks().contains(&c));
p.start_task(b, 2);
p.complete_task(b, 3);
assert!(p.ready_tasks().contains(&c));
}
#[test]
fn default_impl_is_empty() {
let p = StorageMigrationPlanner::default();
assert_eq!(p.stats().total_tasks, 0);
}
}