nexus-core 0.0.1-alpha

use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;

use anyhow::Result;

#[derive(Debug)]
pub struct DeviceState {
    pub id: usize,
    pub path: PathBuf,
    queue_depth: AtomicU64,
    next_offset: AtomicU64,
}

impl DeviceState {
    pub fn new(id: usize, path: PathBuf) -> Self {
        Self {
            id,
            path,
            queue_depth: AtomicU64::new(0),
            next_offset: AtomicU64::new(0),
        }
    }

    pub fn queue_depth(&self) -> u64 {
        self.queue_depth.load(Ordering::Relaxed)
    }

    pub fn reserve_aligned_offset(&self, len: u64, alignment: u64) -> u64 {
        let aligned_len = align_up_u64(len, alignment);
        self.next_offset.fetch_add(aligned_len, Ordering::AcqRel)
    }

    pub fn increment_queue_depth(&self) {
        self.queue_depth.fetch_add(1, Ordering::AcqRel);
    }

    pub fn decrement_queue_depth(&self) {
        self.queue_depth
            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |depth| {
                Some(depth.saturating_sub(1))
            })
            .ok();
    }
}

#[derive(Debug, Clone)]
pub struct DeviceScheduler {
    devices: Vec<Arc<DeviceState>>,
    rr_counter: Arc<AtomicUsize>,
}

#[derive(Debug, Clone)]
pub struct ScheduledDevice {
    pub device: Arc<DeviceState>,
    pub reserved_offset: u64,
}

impl DeviceScheduler {
    pub fn new(device_paths: &[PathBuf]) -> Result<Self> {
        if device_paths.is_empty() {
            anyhow::bail!("at least one device path is required");
        }
        if device_paths.len() > 5 {
            anyhow::bail!("at most 5 data devices are supported");
        }

        let devices = device_paths
            .iter()
            .enumerate()
            .map(|(idx, path)| Arc::new(DeviceState::new(idx, path.clone())))
            .collect::<Vec<_>>();

        Ok(Self {
            devices,
            rr_counter: Arc::new(AtomicUsize::new(0)),
        })
    }

    pub fn devices(&self) -> &[Arc<DeviceState>] {
        &self.devices
    }

    pub fn select_and_reserve(&self, len: u64, alignment: u64) -> ScheduledDevice {
        let idx = self.select_device_index();
        let device = self.devices[idx].clone();
        let reserved_offset = device.reserve_aligned_offset(len, alignment);
        ScheduledDevice {
            device,
            reserved_offset,
        }
    }

    pub fn telemetry(&self) -> Vec<(usize, PathBuf, u64)> {
        self.devices
            .iter()
            .map(|device| (device.id, device.path.clone(), device.queue_depth()))
            .collect()
    }

    fn select_device_index(&self) -> usize {
        if self.devices.len() == 1 {
            return 0;
        }

        let mut min_depth = u64::MAX;
        let mut candidates = Vec::<usize>::new();

        for (idx, device) in self.devices.iter().enumerate() {
            let depth = device.queue_depth();
            match depth.cmp(&min_depth) {
                std::cmp::Ordering::Less => {
                    min_depth = depth;
                    candidates.clear();
                    candidates.push(idx);
                }
                std::cmp::Ordering::Equal => {
                    candidates.push(idx);
                }
                std::cmp::Ordering::Greater => {}
            }
        }

        if candidates.len() == 1 {
            return candidates[0];
        }

        let start = self.rr_counter.fetch_add(1, Ordering::AcqRel) % self.devices.len();
        for offset in 0..self.devices.len() {
            let idx = (start + offset) % self.devices.len();
            if candidates.contains(&idx) {
                return idx;
            }
        }
        candidates[0]
    }
}

pub fn align_up_u64(value: u64, alignment: u64) -> u64 {
    if value % alignment == 0 {
        value
    } else {
        value + (alignment - (value % alignment))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn select_and_reserve_never_overlaps_offsets() {
        let paths = vec![
            PathBuf::from("/dev/device-a"),
            PathBuf::from("/dev/device-b"),
        ];
        let scheduler = DeviceScheduler::new(&paths).expect("scheduler init should succeed");
        let first = scheduler.select_and_reserve(4096, 4096);
        let second = scheduler.select_and_reserve(4096, 4096);
        if first.device.id == second.device.id {
            assert_ne!(first.reserved_offset, second.reserved_offset);
        }
    }

    #[test]
    fn align_up_u64_rounds_to_boundary() {
        assert_eq!(align_up_u64(4097, 4096), 8192);
        assert_eq!(align_up_u64(8192, 4096), 8192);
    }

    #[test]
    fn scheduler_rejects_more_than_five_devices() {
        let paths = (0..6)
            .map(|idx| PathBuf::from(format!("/dev/device-{idx}")))
            .collect::<Vec<_>>();
        let err = DeviceScheduler::new(&paths).expect_err("expected validation failure");
        assert!(err.to_string().contains("at most 5 data devices"));
    }
}