yggdrasil/src/device/nvme/mod.rs

#![allow(missing_docs)]
use core::{mem::size_of, time::Duration};

use abi::error::Error;
use alloc::{collections::BTreeMap, vec::Vec};
use device_api::{interrupt::MsiHandler, Device};
use kernel_util::{
    mem::{
        address::{FromRaw, IntoRaw, PhysicalAddress},
        device::{DeviceMemoryIo, DeviceMemoryIoMut},
    },
    sync::IrqSafeSpinlock,
    util::OneTimeInit,
};
use tock_registers::{
    interfaces::{ReadWriteable, Readable, Writeable},
    register_bitfields, register_structs,
    registers::{ReadOnly, ReadWrite, WriteOnly},
};

use crate::{
    arch::{Architecture, ARCHITECTURE},
    device::{
        bus::pci::{
            capability::MsiXCapability, PciBaseAddress, PciCommandRegister, PciConfigurationSpace,
        },
        nvme::{
            command::{
                IdentifyActiveNamespaceIdListRequest, IdentifyControllerRequest, IoRead, IoWrite,
            },
            drive::NvmeDrive,
            queue::{CompletionQueueEntry, SubmissionQueueEntry},
        },
    },
    task::runtime,
};

use self::{
    command::{CreateIoCompletionQueue, CreateIoSubmissionQueue, SetFeatureRequest},
    error::NvmeError,
    queue::QueuePair,
};

use super::bus::pci::{capability::MsiXEntry, FromPciBus, PciDeviceInfo};

mod command;
mod drive;
mod error;
mod queue;

register_bitfields! {
    u32,
    CC [
        IOCQES OFFSET(20) NUMBITS(4) [],
        IOSQES OFFSET(16) NUMBITS(4) [],
        AMS OFFSET(11) NUMBITS(3) [],
        MPS OFFSET(7) NUMBITS(4) [],
        CSS OFFSET(4) NUMBITS(3) [
            NvmCommandSet = 0
        ],
        ENABLE OFFSET(0) NUMBITS(1) [],
    ],
    CSTS [
        CFS OFFSET(1) NUMBITS(1) [],
        RDY OFFSET(0) NUMBITS(1) [],
    ],
    AQA [
        /// Admin Completion Queue Size in entries - 1
        ACQS OFFSET(16) NUMBITS(12) [],
        /// Admin Submission Queue Size in entries - 1
        ASQS OFFSET(0) NUMBITS(12) [],
    ]
}

register_bitfields! {
    u64,
    CAP [
        /// Maximum Queue Entries Supported - 1. i.e., 0 means maximum queue len of 1, 1 = 2 etc.
        MQES OFFSET(0) NUMBITS(16) [],
        /// Timeout. Represents the worst-case time the host software should wait for CSTS.RDY to
        /// change its state.
        TO OFFSET(24) NUMBITS(8) [],
        /// Doorbell stride. Stride in bytes = pow(2, 2 + DSTRD).
        DSTRD OFFSET(32) NUMBITS(4) [],
        /// NVM Subsystem Reset Supported (see NVMe BS Section 3.7.1)
        NSSRS OFFSET(36) NUMBITS(1) [],
        /// Controller supports one or more I/O command sets
        CSS_IO_COMMANDS OFFSET(43) NUMBITS(1) [],
        /// Controller only supports admin commands and no I/O commands
        CSS_ADMIN_ONLY OFFSET(44) NUMBITS(1) [],
        /// Memory page size minimum (bytes = pow(2, 12 + MPSMIN))
        MPSMIN OFFSET(48) NUMBITS(4) [],
        /// Memory page size maximum -|-
        MPSMAX OFFSET(52) NUMBITS(4) [],
    ]
}

register_structs! {
    #[allow(non_snake_case)]
    Regs {
        (0x00 => CAP: ReadOnly<u64, CAP::Register>),
        (0x08 => VS: ReadOnly<u32>),
        (0x0C => INTMS: WriteOnly<u32>),
        (0x10 => INTMC: WriteOnly<u32>),
        (0x14 => CC: ReadWrite<u32, CC::Register>),
        (0x18 => _0),
        (0x1C => CSTS: ReadOnly<u32, CSTS::Register>),
        (0x20 => _1),
        (0x24 => AQA: ReadWrite<u32, AQA::Register>),
        (0x28 => ASQ: ReadWrite<u64>),
        (0x30 => ACQ: ReadWrite<u64>),
        (0x38 => _2),
        (0x2000 => @END),
    }
}

pub struct NvmeController {
    regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,
    admin_q: OneTimeInit<QueuePair>,
    ioqs: OneTimeInit<Vec<QueuePair>>,
    vector_table: IrqSafeSpinlock<DeviceMemoryIoMut<'static, [MsiXEntry]>>,
    drive_table: IrqSafeSpinlock<BTreeMap<u32, &'static NvmeDrive>>,
    controller_id: OneTimeInit<usize>,

    doorbell_shift: usize,
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum IoDirection {
    Read,
    Write,
}

impl Regs {
    unsafe fn doorbell_ptr(&self, shift: usize, completion: bool, queue_index: usize) -> *mut u32 {
        let doorbell_base = (self as *const Regs as *mut Regs).addr() + 0x1000;
        let offset = ((queue_index << shift) + completion as usize) * 4;
        (doorbell_base + offset) as *mut u32
    }
}

impl NvmeController {
    async fn late_init(&'static self) -> Result<(), NvmeError> {
        register_nvme_controller(self);

        let admin_q = self.admin_q.get();

        // Identify the controller
        let _identify = admin_q.request(IdentifyControllerRequest).await?;

        // TODO do something with identify_controller

        // Request a CQ/SQ pair for I/O
        admin_q
            .request_no_data(SetFeatureRequest::NumberOfQueues(1, 1))
            .await?;

        // Allocate the queue
        let (sq_doorbell, cq_doorbell) = unsafe { self.doorbell_pair(1) };
        let io_q =
            QueuePair::new(1, 0, 32, sq_doorbell, cq_doorbell).map_err(NvmeError::MemoryError)?;

        // Create the queue on the device side
        admin_q
            .request_no_data(CreateIoCompletionQueue {
                id: 1,
                size: 32,
                vector: 0,
                data: io_q.cq_physical_pointer(),
            })
            .await?;
        admin_q
            .request_no_data(CreateIoSubmissionQueue {
                id: 1,
                cq_id: 1,
                size: 32,
                data: io_q.sq_physical_pointer(),
            })
            .await?;

        self.ioqs.init(Vec::from_iter([io_q]));

        // Identify namespaces
        self.enumerate_namespaces().await?;

        Ok(())
    }

    async fn enumerate_namespaces(&'static self) -> Result<(), NvmeError> {
        let admin_q = self.admin_q.get();

        let namespaces = admin_q
            .request(IdentifyActiveNamespaceIdListRequest { start_id: 0 })
            .await?;

        let count = namespaces.entries.iter().position(|&x| x == 0).unwrap();
        let list = &namespaces.entries[..count];

        for &nsid in list {
            match NvmeDrive::create(self, nsid).await {
                Ok(drive) => {
                    self.drive_table.lock().insert(nsid, drive);
                }
                Err(error) => {
                    warnln!("Could not create nvme drive, nsid={}: {:?}", nsid, error);
                }
            }
        }

        Ok(())
    }

    // TODO sane methods for IO
    pub async fn perform_io(
        &'static self,
        nsid: u32,
        lba: u64,
        buffer_address: PhysicalAddress,
        direction: IoDirection,
    ) -> Result<(), NvmeError> {
        let ioq = &self.ioqs.get()[0];

        debugln!("{:?} nsid={}, lba={:#x}", direction, nsid, lba);
        let cmd_id = match direction {
            IoDirection::Read => ioq.submit(
                IoRead {
                    nsid,
                    lba,
                    count: 1,
                },
                &[buffer_address],
                true,
            ),
            IoDirection::Write => ioq.submit(
                IoWrite {
                    nsid,
                    lba,
                    count: 1,
                },
                &[buffer_address],
                true,
            ),
        };

        ioq.wait_for_completion(cmd_id, ()).await?;

        Ok(())
    }

    unsafe fn doorbell_pair(&self, idx: usize) -> (*mut u32, *mut u32) {
        let regs = self.regs.lock();
        let sq_ptr = regs.doorbell_ptr(self.doorbell_shift, false, idx);
        let cq_ptr = regs.doorbell_ptr(self.doorbell_shift, true, idx);
        (sq_ptr, cq_ptr)
    }
}

impl MsiHandler for NvmeController {
    fn handle_msi(&self, _vector: usize) -> bool {
        // TODO check MSI-X pending bits
        self.admin_q.get().process_completions();
        if let Some(qs) = self.ioqs.try_get() {
            for q in qs {
                q.process_completions();
            }
        }
        true
    }
}

impl Device for NvmeController {
    unsafe fn init(&'static self) -> Result<(), Error> {
        let regs = self.regs.lock();

        let min_page_size = 1usize << (12 + regs.CAP.read(CAP::MPSMIN));

        if min_page_size > 4096 {
            panic!();
        }

        let timeout = Duration::from_millis(regs.CAP.read(CAP::TO) * 500);
        debugln!("Worst-case timeout: {:?}", timeout);

        while regs.CSTS.matches_any(CSTS::RDY::SET) {
            core::hint::spin_loop();
        }

        let queue_slots = 32;
        if queue_slots > regs.CAP.read(CAP::MQES) + 1 {
            todo!(
                "queue_slots too big, max = {}",
                regs.CAP.read(CAP::MQES) + 1
            );
        }

        // Setup the admin queue (index 0)
        let admin_sq_doorbell = unsafe { regs.doorbell_ptr(self.doorbell_shift, false, 0) };
        let admin_cq_doorbell = unsafe { regs.doorbell_ptr(self.doorbell_shift, true, 0) };
        debugln!("sq_doorbell for adminq = {:p}", admin_sq_doorbell);
        let admin_q = QueuePair::new(
            0,
            0,
            queue_slots as usize,
            admin_sq_doorbell,
            admin_cq_doorbell,
        )
        .unwrap();

        regs.AQA
            .modify(AQA::ASQS.val(queue_slots as u32 - 1) + AQA::ACQS.val(queue_slots as u32 - 1));
        regs.ASQ.set(admin_q.sq_physical_pointer().into_raw());
        regs.ACQ.set(admin_q.cq_physical_pointer().into_raw());

        // Configure the controller
        const IOSQES: u32 = size_of::<SubmissionQueueEntry>().ilog2();
        const IOCQES: u32 = size_of::<CompletionQueueEntry>().ilog2();

        regs.CC.modify(
            CC::IOCQES.val(IOCQES)
                + CC::IOSQES.val(IOSQES)
                + CC::MPS.val(0)
                + CC::CSS::NvmCommandSet,
        );

        // Enable the controller
        regs.CC.modify(CC::ENABLE::SET);

        debugln!("Reset the controller");

        while !regs.CSTS.matches_any(CSTS::RDY::SET + CSTS::CFS::SET) {
            core::hint::spin_loop();
        }

        if regs.CSTS.matches_any(CSTS::CFS::SET) {
            todo!("CFS set after reset!");
        }

        self.admin_q.init(admin_q);

        // Register the IRQs (TODO: use multiple)
        {
            let mut vt = self.vector_table.lock();

            // Register vector 0
            vt[0]
                .register(ARCHITECTURE.message_interrupt_controller(), self)
                .unwrap();
        }

        // Schedule late_init task
        runtime::spawn(self.late_init())?;

        Ok(())
    }

    fn display_name(&self) -> &'static str {
        "NVM Express Controller"
    }
}

impl FromPciBus for NvmeController {
    fn from_pci_bus(info: &PciDeviceInfo) -> Result<Self, Error> {
        let PciBaseAddress::Memory(bar0) = info.config_space.bar(0).unwrap() else {
            panic!();
        };

        // TODO also support MSI
        let mut msix = info.config_space.capability::<MsiXCapability>().unwrap();
        let mut vt = msix.vector_table()?;

        for vector in vt.iter_mut() {
            vector.set_masked(true);
        }
        msix.set_enabled(true);

        let mut cmd = PciCommandRegister::from_bits_retain(info.config_space.command());
        cmd &= !(PciCommandRegister::DISABLE_INTERRUPTS | PciCommandRegister::ENABLE_IO);
        cmd |= PciCommandRegister::ENABLE_MEMORY | PciCommandRegister::BUS_MASTER;
        info.config_space.set_command(cmd.bits());

        let regs = unsafe { DeviceMemoryIo::<Regs>::map(PhysicalAddress::from_raw(bar0)) }?;

        // Disable the controller
        regs.CC.modify(CC::ENABLE::CLEAR);

        let doorbell_shift = regs.CAP.read(CAP::DSTRD) as usize + 1;

        Ok(Self {
            regs: IrqSafeSpinlock::new(regs),
            admin_q: OneTimeInit::new(),
            ioqs: OneTimeInit::new(),
            vector_table: IrqSafeSpinlock::new(vt),
            drive_table: IrqSafeSpinlock::new(BTreeMap::new()),
            controller_id: OneTimeInit::new(),
            doorbell_shift,
        })
    }
}

static NVME_CONTROLLERS: IrqSafeSpinlock<Vec<&'static NvmeController>> =
    IrqSafeSpinlock::new(Vec::new());

pub fn register_nvme_controller(ctrl: &'static NvmeController) {
    let mut list = NVME_CONTROLLERS.lock();
    let id = list.len();
    list.push(ctrl);
    ctrl.controller_id.init(id);
}