use super::config::DriveConfig; use anyhow::Result; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use std::path::{Path, PathBuf}; use tokio::fs; // ============================ // Drive format (on-disk metadata) // ============================ #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct DriveFormat { pub cluster_id: String, pub erasure_set_id: u32, pub drive_index_in_set: u32, pub format_version: u32, } // ============================ // Drive state tracking // ============================ #[derive(Debug, Clone, PartialEq, Eq)] pub enum DriveStatus { Online, Degraded, Offline, Healing, } #[derive(Debug, Clone)] pub struct DriveStats { pub total_bytes: u64, pub available_bytes: u64, pub used_bytes: u64, pub avg_write_latency_us: u64, pub avg_read_latency_us: u64, pub error_count: u64, pub last_error: Option, pub last_check: DateTime, } impl Default for DriveStats { fn default() -> Self { Self { total_bytes: 0, available_bytes: 0, used_bytes: 0, avg_write_latency_us: 0, avg_read_latency_us: 0, error_count: 0, last_error: None, last_check: Utc::now(), } } } #[derive(Debug, Clone)] pub struct DriveState { pub path: PathBuf, pub format: Option, pub status: DriveStatus, pub stats: DriveStats, } // ============================ // Drive manager // ============================ pub struct DriveManager { drives: Vec, } impl DriveManager { /// Initialize drive manager with configured drive paths. pub async fn new(config: &DriveConfig) -> Result { let paths: Vec = config.paths.iter().map(PathBuf::from).collect(); Self::from_paths(&paths).await } /// Initialize drive manager from an explicit list of resolved paths. pub async fn from_paths(paths: &[PathBuf]) -> Result { let mut drives = Vec::with_capacity(paths.len()); for path in paths { let storage_dir = path.join(".smartstorage"); // Ensure the drive directory exists fs::create_dir_all(&storage_dir).await?; // Try to read existing format let format = Self::read_format(&storage_dir).await; let status = if path.exists() { DriveStatus::Online } else { DriveStatus::Offline }; drives.push(DriveState { path: path.clone(), format, status, stats: DriveStats::default(), }); } Ok(Self { drives }) } /// Format drives for a new cluster. Stamps each drive with cluster and erasure set info. pub async fn format_drives( &mut self, cluster_id: &str, erasure_set_assignments: &[(u32, u32)], // (erasure_set_id, drive_index_in_set) ) -> Result<()> { if erasure_set_assignments.len() != self.drives.len() { anyhow::bail!( "Erasure set assignments count ({}) doesn't match drive count ({})", erasure_set_assignments.len(), self.drives.len() ); } for (drive, (set_id, drive_idx)) in self.drives.iter_mut().zip(erasure_set_assignments.iter()) { let format = DriveFormat { cluster_id: cluster_id.to_string(), erasure_set_id: *set_id, drive_index_in_set: *drive_idx, format_version: 1, }; let storage_dir = drive.path.join(".smartstorage"); fs::create_dir_all(&storage_dir).await?; let format_path = storage_dir.join("format.json"); let json = serde_json::to_string_pretty(&format)?; fs::write(&format_path, json).await?; drive.format = Some(format); } Ok(()) } /// Get the number of drives managed. pub fn drive_count(&self) -> usize { self.drives.len() } /// Get a drive's state by index. pub fn drive(&self, index: usize) -> Option<&DriveState> { self.drives.get(index) } /// Get all drives. pub fn drives(&self) -> &[DriveState] { &self.drives } /// Get a cloneable snapshot of current drive states. pub fn snapshot(&self) -> Vec { self.drives.clone() } /// Get drives that are online. pub fn online_drives(&self) -> Vec { self.drives .iter() .enumerate() .filter(|(_, d)| d.status == DriveStatus::Online) .map(|(i, _)| i) .collect() } /// Check health of a specific drive by writing and reading a probe file. pub async fn check_drive_health(&mut self, index: usize) -> Result { let drive = self .drives .get_mut(index) .ok_or_else(|| anyhow::anyhow!("Drive index {} out of range", index))?; let probe_path = drive.path.join(".smartstorage").join(".health_probe"); let start = std::time::Instant::now(); // Write probe match fs::write(&probe_path, b"health_check").await { Ok(()) => {} Err(e) => { drive.stats.error_count += 1; drive.stats.last_error = Some(e.to_string()); drive.status = DriveStatus::Offline; drive.stats.last_check = Utc::now(); return Ok(DriveStatus::Offline); } } // Read probe match fs::read(&probe_path).await { Ok(_) => {} Err(e) => { drive.stats.error_count += 1; drive.stats.last_error = Some(e.to_string()); drive.status = DriveStatus::Offline; drive.stats.last_check = Utc::now(); return Ok(DriveStatus::Offline); } } // Clean up probe let _ = fs::remove_file(&probe_path).await; let latency = start.elapsed(); if let Some((total_bytes, available_bytes, used_bytes)) = filesystem_usage(&drive.path) { drive.stats.total_bytes = total_bytes; drive.stats.available_bytes = available_bytes; drive.stats.used_bytes = used_bytes; } drive.stats.avg_write_latency_us = latency.as_micros() as u64; drive.stats.last_check = Utc::now(); // Mark degraded if latency is too high (>5 seconds) if latency.as_secs() > 5 { drive.status = DriveStatus::Degraded; } else { drive.status = DriveStatus::Online; } Ok(drive.status.clone()) } /// Run health checks on all drives. pub async fn check_all_drives(&mut self) -> Vec<(usize, DriveStatus)> { let mut results = Vec::new(); let count = self.drives.len(); for i in 0..count { match self.check_drive_health(i).await { Ok(status) => results.push((i, status)), Err(e) => { tracing::error!(drive = i, error = %e, "Drive health check failed"); results.push((i, DriveStatus::Offline)); } } } results } // Internal helpers async fn read_format(storage_dir: &Path) -> Option { let format_path = storage_dir.join("format.json"); let content = fs::read_to_string(&format_path).await.ok()?; serde_json::from_str(&content).ok() } } #[cfg(unix)] fn filesystem_usage(path: &Path) -> Option<(u64, u64, u64)> { use std::ffi::CString; use std::os::unix::ffi::OsStrExt; let path_bytes = path.as_os_str().as_bytes(); let c_path = CString::new(path_bytes).ok()?; let mut stat: libc::statvfs = unsafe { std::mem::zeroed() }; if unsafe { libc::statvfs(c_path.as_ptr(), &mut stat) } != 0 { return None; } let block_size = stat.f_frsize as u64; let total_bytes = stat.f_blocks as u64 * block_size; let available_bytes = stat.f_bavail as u64 * block_size; let free_bytes = stat.f_bfree as u64 * block_size; let used_bytes = total_bytes.saturating_sub(free_bytes); Some((total_bytes, available_bytes, used_bytes)) } #[cfg(not(unix))] fn filesystem_usage(_path: &Path) -> Option<(u64, u64, u64)> { None }