Only trigger future on poll in controller.

rust/sys/denali/src/ahci/controller.rs

@@ -1,5 +1,9 @@
+use core::array;
+use core::future::Future;
+use core::ops::Deref;
 use core::ops::DerefMut;
-use core::ptr::addr_of_mut;
+use core::pin::Pin;
+use core::task::{Context, Poll, Waker};
 
 use alloc::boxed::Box;
 use alloc::sync::Arc;
@@ -73,8 +77,36 @@ impl AhciController {
 
     fn init(&mut self) {
         self.hba.lock().init();
+        let hba = self.hba.lock();
+        for i in 0..(hba.capabilities.read().num_ports() as usize) {
+            let port_index = 1 << i;
+            if (hba.port_implemented.read() & port_index) != port_index {
+                mammoth::debug!("Skipping port {}, not implemented", i);
+                continue;
+            }
 
-        self.init_ports().unwrap();
+            let port_offset: usize = 0x100 + (0x80 * i);
+            let port = unsafe {
+                ((self.hba_vaddr as usize + port_offset) as *mut AhciPortHba)
+                    .as_mut()
+                    .unwrap()
+            };
+
+            let sata_status = port.sata_status.read();
+            if (sata_status.device_detection() != AhciDeviceDetection::CommunicationEstablished)
+                || (sata_status.interface_power_management()
+                    != AhciInterfacePowerManagement::Active)
+            {
+                mammoth::debug!(
+                    "Skipping port {}, no communcation. Status: {:?}",
+                    i,
+                    sata_status
+                );
+                continue;
+            }
+
+            self.ports[i] = Some(PortController::new(port));
+        }
     }
 
     fn irq_num(&self) -> u64 {
@@ -105,37 +137,36 @@ impl AhciController {
         }
     }
 
-    fn init_ports(&mut self) -> Result<(), ZError> {
-        let hba = self.hba.lock();
-        for i in 0..(hba.capabilities.read().num_ports() as usize) {
-            let port_index = 1 << i;
-            if (hba.port_implemented.read() & port_index) != port_index {
-                mammoth::debug!("Skipping port {}, not implemented", i);
-                continue;
-            }
-
-            let port_offset: usize = 0x100 + (0x80 * i);
-            let port = unsafe {
-                ((self.hba_vaddr as usize + port_offset) as *mut AhciPortHba)
-                    .as_mut()
-                    .unwrap()
+    pub async fn identify_ports(&self) -> Result<(), ZError> {
+        for port in self.ports.iter().flatten() {
+            let sig = port.get_signature();
+            if sig == 0x101 {
+                let command = Command::identify()?;
+                mammoth::debug!("IDENT!");
+                port.issue_command(&command)?.await;
+                let ident = command.memory_region.slice::<u16>();
+                let new_sector_size = if ident[106] & (1 << 12) != 0 {
+                    ident[117] as u32 | ((ident[118] as u32) << 16)
+                } else {
+                    512
                 };
 
-            let sata_status = port.sata_status.read();
-            if (sata_status.device_detection() != AhciDeviceDetection::CommunicationEstablished)
-                || (sata_status.interface_power_management()
-                    != AhciInterfacePowerManagement::Active)
-            {
-                mammoth::debug!(
-                    "Skipping port {}, no communcation. Status: {:?}",
-                    i,
-                    sata_status
-                );
-                continue;
-            }
+                let lba_count = if ident[83] & (1 << 10) != 0 {
+                    ident[100] as u64
+                        | (ident[101] as u64) << 16
+                        | (ident[102] as u64) << 32
+                        | (ident[103] as u64) << 48
+                } else {
+                    ident[60] as u64 | (ident[61] as u64) << 16
+                };
+                mammoth::debug!("Sector size: {:#0x}", new_sector_size);
+                mammoth::debug!("LBA Count: {:#0x}", lba_count);
 
-            self.ports[i] = Some(PortController::new(port));
-            self.ports[i].as_mut().unwrap().identify()?;
+                //self.sector_size = Some(new_sector_size as u64);
+                //self.sector_cnt = Some(lba_count);
+            } else {
+                mammoth::debug!("Skipping non-sata sig: {:#0x}", sig);
+            }
         }
         Ok(())
     }
@@ -151,12 +182,60 @@ pub fn spawn_irq_thread(controller: Arc<AhciController>) -> thread::JoinHandle {
         });
         loop {
             irq_port.recv_null().unwrap();
+            mammoth::debug!("Interrupt!");
             controller.handle_irq();
         }
     };
     thread::spawn(irq_thread)
 }
 
+pub async fn identify_ports(controller: Arc<AhciController>) {
+    controller.identify_ports().await.unwrap();
+}
+
+enum CommandStatus {
+    Empty,
+    NotSent,
+    Pending(Waker),
+    Complete,
+}
+
+struct CommandFuture {
+    status: Arc<Mutex<CommandStatus>>,
+    trigger: Box<dyn Fn() + Sync>,
+}
+
+impl CommandFuture {
+    fn new(status: Arc<Mutex<CommandStatus>>, trigger: Box<dyn Fn() + Sync>) -> Self {
+        Self { status, trigger }
+    }
+}
+
+impl Drop for CommandFuture {
+    fn drop(&mut self) {
+        *self.status.lock() = CommandStatus::Empty;
+    }
+}
+
+impl Future for CommandFuture {
+    type Output = ();
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let s = self.deref_mut();
+        let mut status = s.status.lock();
+        match status.deref() {
+            CommandStatus::NotSent => {
+                *status = CommandStatus::Pending(cx.waker().clone());
+                (s.trigger)();
+                Poll::Pending
+            }
+            CommandStatus::Pending(_) => Poll::Pending,
+            CommandStatus::Complete => Poll::Ready(()),
+            CommandStatus::Empty => panic!("Polling empty command slot"),
+        }
+    }
+}
+
 struct Command {
     command: SataCommand,
     lba: u64,
@@ -165,12 +244,10 @@ struct Command {
 
     #[allow(dead_code)] // We need to own this even if we never access it.
     memory_region: MemoryRegion,
-
-    callback_internal: Box<dyn Fn(&Self) + Sync + Send + 'static>,
 }
 
 impl Command {
-    pub fn identify(callback: Box<dyn Fn(&Self) + Sync + Send + 'static>) -> Result<Self, ZError> {
+    pub fn identify() -> Result<Self, ZError> {
         let (memory_region, paddr) = MemoryRegion::contiguous_physical(512)?;
 
         Ok(Self {
@@ -179,13 +256,8 @@ impl Command {
             sector_cnt: 1,
             paddr,
             memory_region,
-            callback_internal: callback,
         })
     }
-
-    pub fn callback(&self) {
-        (self.callback_internal)(self);
-    }
 }
 
 impl From<&Command> for HostToDeviceRegisterFis {
@@ -194,17 +266,17 @@ impl From<&Command> for HostToDeviceRegisterFis {
     }
 }
 
-struct PortController {
-    ahci_port_hba: Mutex<&'static mut AhciPortHba>,
+struct CommandStructures {
     command_list: &'static mut CommandList,
     received_fis: &'static mut ReceivedFis,
     command_tables: &'static mut [CommandTable; 32],
+}
 
-    command_slots: Mutex<[Option<Arc<Command>>; 32]>,
+struct PortController {
+    ahci_port_hba: Arc<Mutex<&'static mut AhciPortHba>>,
+    command_structures: Mutex<CommandStructures>,
 
-    // FIXME: These should probably be something like a OnceCell (or OnceLock).
-    sector_size: Arc<Mutex<Option<u64>>>,
-    sector_cnt: Arc<Mutex<Option<u64>>>,
+    command_slots: [Arc<Mutex<CommandStatus>>; 32],
 }
 
 impl PortController {
@@ -230,76 +302,61 @@ impl PortController {
                 (command_paddr + 0x500) + (0x100 * (i as u64));
         }
 
+        let command_slots = array::from_fn(|_| Arc::new(Mutex::new(CommandStatus::Empty)));
+
         Self {
-            ahci_port_hba: Mutex::new(ahci_port_hba),
+            ahci_port_hba: Arc::new(Mutex::new(ahci_port_hba)),
+            command_structures: Mutex::new(CommandStructures {
                 command_list,
                 received_fis,
                 command_tables,
-            command_slots: Mutex::new([const { None }; 32]),
-            sector_size: Arc::new(Mutex::new(None)),
-            sector_cnt: Arc::new(Mutex::new(None)),
+            }),
+            command_slots,
         }
     }
 
-    pub fn identify(&mut self) -> Result<(), ZError> {
-        let sig = self.ahci_port_hba.lock().signature.read();
-        if sig == 0x101 {
-            let sector_size = self.sector_size.clone();
-            let sector_cnt = self.sector_cnt.clone();
-            let callback = move |c: &Command| {
-                mammoth::debug!("TESTING");
-                let ident = c.memory_region.slice::<u16>();
-                let new_sector_size = if ident[106] & (1 << 12) != 0 {
-                    ident[117] as u32 | ((ident[118] as u32) << 16)
-                } else {
-                    512
-                };
-
-                let lba_count = if ident[83] & (1 << 10) != 0 {
-                    ident[100] as u64
-                        | (ident[101] as u64) << 16
-                        | (ident[102] as u64) << 32
-                        | (ident[103] as u64) << 48
-                } else {
-                    ident[60] as u64 | (ident[61] as u64) << 16
-                };
-                mammoth::debug!("Sector size: {:#0x}", new_sector_size);
-                mammoth::debug!("LBA Count: {:#0x}", lba_count);
-
-                let _ = sector_size
-                    .lock()
-                    .deref_mut()
-                    .insert(new_sector_size as u64);
-                let _ = sector_cnt.lock().deref_mut().insert(lba_count);
-            };
-            self.issue_command(Arc::from(Command::identify(Box::new(callback))?))?;
-        } else {
-            mammoth::debug!("Skipping non-sata sig: {:#0x}", sig);
-        }
-        Ok(())
+    fn get_signature(&self) -> u32 {
+        self.ahci_port_hba.lock().signature.read()
     }
 
-    fn issue_command(&mut self, command: Arc<Command>) -> Result<(), ZError> {
+    fn issue_command(&self, command: &Command) -> Result<CommandFuture, ZError> {
         let slot = self.select_slot()?;
-        self.command_slots.lock()[slot] = Some(command.clone());
+        let command_slot = self.command_slots[slot].clone();
 
-        self.command_tables[slot].command_fis.host_to_device = command.clone().as_ref().into();
+        let ahci_port_hba_clone = self.ahci_port_hba.clone();
+        let slot_clone = slot;
+        let future = CommandFuture::new(
+            command_slot,
+            Box::new(move || {
+                ahci_port_hba_clone.lock().issue_command(slot_clone);
+            }),
+        );
 
-        self.command_tables[slot].prdt[0].region_address = command.paddr;
-        self.command_tables[slot].prdt[0].byte_count = 512 * (command.sector_cnt as u32);
+        let mut command_structures = self.command_structures.lock();
 
-        self.command_list[slot].prd_table_length = 1;
+        command_structures.command_tables[slot]
+            .command_fis
+            .host_to_device = command.into();
 
-        self.command_list[slot].command = (size_of::<HostToDeviceRegisterFis>() as u16 / 4) & 0x1F;
-        self.command_list[slot].command |= 1 << 7;
-        self.ahci_port_hba.lock().issue_command(slot);
-        Ok(())
+        command_structures.command_tables[slot].prdt[0].region_address = command.paddr;
+        command_structures.command_tables[slot].prdt[0].byte_count =
+            512 * (command.sector_cnt as u32);
+
+        command_structures.command_list[slot].prd_table_length = 1;
+
+        command_structures.command_list[slot].command =
+            (size_of::<HostToDeviceRegisterFis>() as u16 / 4) & 0x1F;
+        command_structures.command_list[slot].command |= 1 << 7;
+
+        Ok(future)
     }
 
     fn select_slot(&self) -> Result<usize, ZError> {
-        let slots = self.command_slots.lock();
-        for i in 0..slots.len() {
-            if slots[i].is_none() {
+        // TODO: We have a race condition here.
+        for i in 0..self.command_slots.len() {
+            let mut slot = self.command_slots[i].lock();
+            if matches!(*slot, CommandStatus::Empty) {
+                *slot = CommandStatus::NotSent;
                 return Ok(i);
             }
         }
@@ -309,20 +366,18 @@ impl PortController {
     fn handle_interrupt(&self) {
         let int_status = self.ahci_port_hba.lock().interrupt_status.read();
         if int_status.device_to_host_register_fis_interrupt() {
+            let received_fis = self
+                .command_structures
+                .lock()
+                .received_fis
+                .device_to_host_register_fis;
             assert_eq!(
-                self.received_fis.device_to_host_register_fis.fis_type as u8,
+                received_fis.fis_type as u8,
                 FisType::RegisterDeviceToHost as u8
             );
-            if self.received_fis.device_to_host_register_fis.error != 0 {
-                mammoth::debug!(
-                    "D2H err: {:#0x}",
-                    self.received_fis.device_to_host_register_fis.error
-                );
-
-                mammoth::debug!(
-                    "Status: {:#0x}",
-                    self.received_fis.device_to_host_register_fis.status
-                );
+            if received_fis.error != 0 {
+                mammoth::debug!("D2H err: {:#0x}", received_fis.error);
+                mammoth::debug!("Status: {:#0x}", received_fis.status);
             }
 
             self.ahci_port_hba.lock().interrupt_status.write(
@@ -340,20 +395,19 @@ impl PortController {
             let int_offset = 1 << i;
 
             // If there is no longer a command issued on a slot and we have something in
-            // the command list we know that this is the command that finished.
-            // FIXME: This could cause a race condition when issuing a command if a different
-            // interrupt triggers between us setting the command in the command slot and
-            // actually issuing the command.
-            if (self.ahci_port_hba.lock().command_issue.read() & int_offset) != int_offset
-                && self.command_slots.lock()[i].is_some()
-            {
-                self.finish_command(i);
-                self.command_slots.lock()[i] = None;
-            }
-        }
+            // the command list that is pending we know that this is the command that finished.
+            if (self.ahci_port_hba.lock().command_issue.read() & int_offset) != int_offset {
+                let mut command_status = self.command_slots[i].lock();
+                let mut did_complete = false;
+                if let CommandStatus::Pending(ref waker) = command_status.deref() {
+                    waker.wake_by_ref();
+                    did_complete = true;
                 }
 
-    fn finish_command(&self, slot: usize) {
-        self.command_slots.lock()[slot].as_ref().unwrap().callback()
+                if did_complete {
+                    *command_status = CommandStatus::Complete;
+                }
+            }
+        }
     }
 }

rust/sys/denali/src/ahci/mod.rs

@@ -3,5 +3,6 @@ mod controller;
 mod hba;
 mod port;
 
+pub use controller::identify_ports;
 pub use controller::spawn_irq_thread;
 pub use controller::AhciController;

rust/sys/denali/src/bin/denali.rs

@@ -4,9 +4,14 @@
 extern crate alloc;
 
 use alloc::sync::Arc;
-use mammoth::{define_entry, sync::Mutex, zion::z_err_t};
+use mammoth::{
+    define_entry,
+    sync::Mutex,
+    task::{Executor, Task},
+    zion::z_err_t,
+};
 
-use denali::ahci::{spawn_irq_thread, AhciController};
+use denali::ahci::{identify_ports, spawn_irq_thread, AhciController};
 
 define_entry!();
 
@@ -24,8 +29,14 @@ extern "C" fn main() -> z_err_t {
         ahci_info.ahci_region,
     )));
 
+    let mut executor = Executor::new();
+
+    executor.spawn(Task::new(identify_ports(ahci_controller.clone())));
+
     let thread = spawn_irq_thread(ahci_controller.clone());
 
+    executor.run();
+
     thread.join().expect("Failed to wait on irq thread.");
     0
 }