markbase/markbase-raid/src/raid/level_5.rs

use super::controller::RaidArray;
use super::parity::calculate_new_parity;
use super::{MemberStatus, RaidAlgorithm, RaidError, RaidLevel};
use std::collections::HashMap;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom, Write};
use std::sync::Arc;

pub struct Raid5 {
    array: Arc<RaidArray>,
    stripe_size: u64,
    member_files: HashMap<usize, File>,
}

impl Raid5 {
    pub fn new(array: Arc<RaidArray>) -> Result<Self, RaidError> {
        if array.members.len() < 3 {
            return Err("RAID 5 requires at least 3 disks".into());
        }

        let stripe_size = array.stripe_size;
        let mut member_files = HashMap::new();

        for (i, member) in array.members.iter().enumerate() {
            let file = File::options()
                .read(true)
                .write(true)
                .create(false)
                .open(&member.device_path)?;
            member_files.insert(i, file);
        }

        Ok(Raid5 {
            array,
            stripe_size,
            member_files,
        })
    }

    fn locate_stripe(&self, block_offset: u64) -> (usize, usize, u64) {
        let total_data_disks = self.array.members.len() - 1;
        let stripe_index = (block_offset / self.stripe_size) as usize;
        let offset_in_stripe = block_offset % self.stripe_size;

        let parity_disk = stripe_index % self.array.members.len();
        let data_disk_index = stripe_index % total_data_disks;

        let data_disk = if data_disk_index < parity_disk {
            data_disk_index
        } else {
            data_disk_index + 1
        };

        let physical_offset =
            (stripe_index / total_data_disks) as u64 * self.stripe_size + offset_in_stripe;

        (data_disk, parity_disk, physical_offset)
    }

    fn read_from_member(
        &mut self,
        member_index: usize,
        offset: u64,
        size: u64,
    ) -> Result<Vec<u8>, RaidError> {
        if self.array.members[member_index].status != MemberStatus::Online {
            return Err(format!("Member {} is offline", member_index).into());
        }

        let file = self
            .member_files
            .get_mut(&member_index)
            .ok_or("Member file not found")?;

        file.seek(SeekFrom::Start(offset))?;
        let mut buffer = vec![0u8; size as usize];
        file.read_exact(&mut buffer)?;

        Ok(buffer)
    }

    fn write_to_member(
        &mut self,
        member_index: usize,
        offset: u64,
        data: &[u8],
    ) -> Result<(), RaidError> {
        if self.array.members[member_index].status != MemberStatus::Online {
            return Err(format!("Member {} is offline", member_index).into());
        }

        let file = self
            .member_files
            .get_mut(&member_index)
            .ok_or("Member file not found")?;

        file.seek(SeekFrom::Start(offset))?;
        file.write_all(data)?;

        Ok(())
    }
}

impl RaidAlgorithm for Raid5 {
    fn read(&mut self, block_offset: u64, size: u64) -> Result<Vec<u8>, RaidError> {
        let mut result = Vec::with_capacity(size as usize);
        let mut remaining = size;
        let mut current_offset = block_offset;

        while remaining > 0 {
            let (data_disk, _parity_disk, physical_offset) = self.locate_stripe(current_offset);
            let chunk_size = std::cmp::min(
                remaining,
                self.stripe_size - (current_offset % self.stripe_size),
            );

            let data = self.read_from_member(data_disk, physical_offset, chunk_size)?;
            result.extend_from_slice(&data);

            remaining -= chunk_size;
            current_offset += chunk_size;
        }

        Ok(result)
    }

    fn write(&mut self, block_offset: u64, data: &[u8]) -> Result<(), RaidError> {
        let mut remaining = data.len() as u64;
        let mut current_offset = block_offset;
        let mut data_pos = 0;

        while remaining > 0 {
            let (data_disk, parity_disk, physical_offset) = self.locate_stripe(current_offset);
            let chunk_size = std::cmp::min(
                remaining,
                self.stripe_size - (current_offset % self.stripe_size),
            );

            let chunk_data = &data[data_pos as usize..(data_pos + chunk_size as usize) as usize];

            let old_data = self.read_from_member(data_disk, physical_offset, chunk_size)?;
            let old_parity = self.read_from_member(parity_disk, physical_offset, chunk_size)?;

            let new_parity = calculate_new_parity(&old_parity, &old_data, chunk_data);

            self.write_to_member(data_disk, physical_offset, chunk_data)?;
            self.write_to_member(parity_disk, physical_offset, &new_parity)?;

            remaining -= chunk_size;
            current_offset += chunk_size;
            data_pos += chunk_size as usize;
        }

        Ok(())
    }

    fn get_total_size(&self) -> u64 {
        self.array.total_size
    }

    fn get_level(&self) -> RaidLevel {
        RaidLevel::RAID5
    }
}

#[cfg(test)]
mod tests {
    use super::super::controller::{RaidArray, RaidMember};
    use super::*;
    use std::path::PathBuf;

    #[test]
    fn test_raid5_stripe_location_logic() {
        let members = vec![
            RaidMember {
                device_id: "member_0".to_string(),
                device_path: PathBuf::from("/tmp/disk0"),
                size: 1024,
                status: MemberStatus::Online,
            },
            RaidMember {
                device_id: "member_1".to_string(),
                device_path: PathBuf::from("/tmp/disk1"),
                size: 1024,
                status: MemberStatus::Online,
            },
            RaidMember {
                device_id: "member_2".to_string(),
                device_path: PathBuf::from("/tmp/disk2"),
                size: 1024,
                status: MemberStatus::Online,
            },
        ];

        let array = Arc::new(RaidArray {
            raid_level: RaidLevel::RAID5,
            members,
            stripe_size: 64 * 1024,
            total_size: 2 * 1024 * 1024,
        });

        let raid5 = Raid5 {
            array,
            stripe_size: 64 * 1024,
            member_files: HashMap::new(),
        };

        let (data_disk, parity_disk, offset) = raid5.locate_stripe(0);
        assert_eq!(parity_disk, 0);
        assert_eq!(data_disk, 1);
        assert_eq!(offset, 0);

        let (data_disk, parity_disk, _) = raid5.locate_stripe(64 * 1024);
        assert_eq!(parity_disk, 1);
        assert!(data_disk != 1);
    }
}