From b270c7c9aa2035f98448f978cadab2a5db6d0c6f Mon Sep 17 00:00:00 2001
From: Drew Galbraith <dgalbraith33@gmail.com>
Date: Sat, 1 Feb 2025 01:37:51 -0800
Subject: [PATCH] Only trigger future on poll in controller.

---
 rust/sys/denali/src/ahci/controller.rs | 40 ++++++++++++++------------
 1 file changed, 21 insertions(+), 19 deletions(-)

diff --git a/rust/sys/denali/src/ahci/controller.rs b/rust/sys/denali/src/ahci/controller.rs
index 5c6849a..4453e13 100644
--- a/rust/sys/denali/src/ahci/controller.rs
+++ b/rust/sys/denali/src/ahci/controller.rs
@@ -5,6 +5,7 @@ use core::ops::DerefMut;
 use core::pin::Pin;
 use core::task::{Context, Poll, Waker};
 
+use alloc::boxed::Box;
 use alloc::sync::Arc;
 use mammoth::cap::Capability;
 use mammoth::sync::Mutex;
@@ -201,11 +202,12 @@ enum CommandStatus {
 
 struct CommandFuture {
     status: Arc<Mutex<CommandStatus>>,
+    trigger: Box<dyn Fn() + Sync>,
 }
 
 impl CommandFuture {
-    fn new(status: Arc<Mutex<CommandStatus>>) -> Self {
-        Self { status }
+    fn new(status: Arc<Mutex<CommandStatus>>, trigger: Box<dyn Fn() + Sync>) -> Self {
+        Self { status, trigger }
     }
 }
 
@@ -219,10 +221,12 @@ impl Future for CommandFuture {
     type Output = ();
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let mut status = self.deref_mut().status.lock();
+        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,
@@ -269,7 +273,7 @@ struct CommandStructures {
 }
 
 struct PortController {
-    ahci_port_hba: Mutex<&'static mut AhciPortHba>,
+    ahci_port_hba: Arc<Mutex<&'static mut AhciPortHba>>,
     command_structures: Mutex<CommandStructures>,
 
     command_slots: [Arc<Mutex<CommandStatus>>; 32],
@@ -301,7 +305,7 @@ impl PortController {
         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,
@@ -318,7 +322,15 @@ impl PortController {
     fn issue_command(&self, command: &Command) -> Result<CommandFuture, ZError> {
         let slot = self.select_slot()?;
         let command_slot = self.command_slots[slot].clone();
-        let future = CommandFuture::new(command_slot);
+
+        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);
+            }),
+        );
 
         let mut command_structures = self.command_structures.lock();
 
@@ -336,10 +348,6 @@ impl PortController {
             (size_of::<HostToDeviceRegisterFis>() as u16 / 4) & 0x1F;
         command_structures.command_list[slot].command |= 1 << 7;
 
-        // FIXME: This is technically a poor future implementation since it starts work before it
-        // is polled.
-        self.ahci_port_hba.lock().issue_command(slot);
-
         Ok(future)
     }
 
@@ -391,15 +399,9 @@ impl PortController {
             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;
-                match command_status.deref() {
-                    CommandStatus::NotSent => {
-                        did_complete = true;
-                    }
-                    CommandStatus::Pending(ref waker) => {
-                        waker.wake_by_ref();
-                        did_complete = true;
-                    }
-                    _ => {}
+                if let CommandStatus::Pending(ref waker) = command_status.deref() {
+                    waker.wake_by_ref();
+                    did_complete = true;
                 }
 
                 if did_complete {