Skip programming erased pages

Flash memory is set to all `1` when erased. We have to erase pages
before writing so writing pages full of `1` is not a good use of
CPU time.

Before
real    2m0.817s
user    0m4.380s
sys     0m3.305s

After
real	1m44.543s
user	0m3.586s
sys	0m2.710s

Author: labbott

drv/cosmo-hf/src/apob.rs

@@ -8,7 +8,8 @@
 //! to Rev. 2.0 February 2025.
 
 use crate::{
-    hf::ServerImpl, FlashAddr, FlashDriver, PAGE_SIZE_BYTES, SECTOR_SIZE_BYTES,
+    hf::HfBufs, hf::ServerImpl, FlashAddr, FlashDriver, PAGE_SIZE_BYTES,
+    SECTOR_SIZE_BYTES,
 };
 use drv_hf_api::{
     ApobBeginError, ApobCommitError, ApobHash, ApobReadError, ApobWriteError,
@@ -237,34 +238,6 @@ impl ApobSlot {
     }
 }
 
-pub(crate) struct ApobBufs {
-    persistent_data: &'static mut [u8; APOB_PERSISTENT_DATA_STRIDE],
-    page: &'static mut [u8; PAGE_SIZE_BYTES],
-    scratch: &'static mut [u8; PAGE_SIZE_BYTES],
-}
-
-/// Grabs references to the static buffers.  Can only be called once.
-impl ApobBufs {
-    pub fn claim_statics() -> Self {
-        use static_cell::ClaimOnceCell;
-        static BUFS: ClaimOnceCell<(
-            [u8; APOB_PERSISTENT_DATA_STRIDE],
-            [u8; PAGE_SIZE_BYTES],
-            [u8; PAGE_SIZE_BYTES],
-        )> = ClaimOnceCell::new((
-            [0; APOB_PERSISTENT_DATA_STRIDE],
-            [0; PAGE_SIZE_BYTES],
-            [0; PAGE_SIZE_BYTES],
-        ));
-        let (persistent_data, page, scratch) = BUFS.claim();
-        Self {
-            persistent_data,
-            page,
-            scratch,
-        }
-    }
-}
-
 /// State machine data, which implements the logic from RFD 593
 ///
 /// See rfd.shared.oxide.computer/rfd/593#_production_strength_implementation
@@ -397,7 +370,7 @@ impl ApobState {
     ///
     /// Searches for an active slot in the metadata regions, updating the offset
     /// in the FPGA driver if found, and erases unused or invalid slots.
-    pub(crate) fn init(drv: &mut FlashDriver, buf: &mut ApobBufs) -> Self {
+    pub(crate) fn init(drv: &mut FlashDriver, buf: &mut HfBufs) -> Self {
         // Look up persistent data, which specifies an active slot
         let out = if let Some(s) = Self::get_slot(drv) {
             // Erase the inactive slot, in preparation for writing
@@ -448,7 +421,7 @@ impl ApobState {
     }
 
     /// Erases the given APOB slot
-    fn slot_erase(drv: &mut FlashDriver, buf: &mut ApobBufs, slot: ApobSlot) {
+    fn slot_erase(drv: &mut FlashDriver, buf: &mut HfBufs, slot: ApobSlot) {
         static_assertions::const_assert!(
             APOB_SLOT_SIZE.is_multiple_of(SECTOR_SIZE_BYTES)
         );
@@ -463,7 +436,7 @@ impl ApobState {
     /// If `size > APOB_SLOT_SIZE`
     fn slot_erase_range(
         drv: &mut FlashDriver,
-        buf: &mut ApobBufs,
+        buf: &mut HfBufs,
         slot: ApobSlot,
         size: u32,
     ) {
@@ -578,7 +551,7 @@ impl ApobState {
     pub(crate) fn write(
         &mut self,
         drv: &mut FlashDriver,
-        buf: &mut ApobBufs,
+        buf: &mut HfBufs,
         offset: u32,
         data: Leased<R, [u8]>,
     ) -> Result<(), ApobWriteError> {
@@ -647,7 +620,7 @@ impl ApobState {
             // If any byte is not a match, then we have to do a flash write
             // (otherwise, it's an idempotent write and we can skip it)
             if needs_write {
-                drv.flash_write(addr, &mut &buf.page[..n])
+                drv.flash_write(addr, &buf.page[..n])
                     .map_err(|_| ApobWriteError::WriteFailed)?;
             }
         }
@@ -657,7 +630,7 @@ impl ApobState {
     pub(crate) fn read(
         &mut self,
         drv: &mut FlashDriver,
-        buf: &mut ApobBufs,
+        buf: &mut HfBufs,
         offset: u32,
         data: Leased<W, [u8]>,
     ) -> Result<usize, ApobReadError> {
@@ -716,7 +689,7 @@ impl ApobState {
     pub(crate) fn commit(
         &mut self,
         drv: &mut FlashDriver,
-        buf: &mut ApobBufs,
+        buf: &mut HfBufs,
     ) -> Result<(), ApobCommitError> {
         drv.check_flash_mux_state()
             .map_err(|_| ApobCommitError::InvalidState)?;
@@ -776,7 +749,7 @@ impl ApobState {
 
     fn apob_validate(
         drv: &mut FlashDriver,
-        buf: &mut ApobBufs,
+        buf: &mut HfBufs,
         expected_length: u32,
         expected_hash: ApobHash,
         write_slot: ApobSlot,
@@ -852,17 +825,17 @@ impl ApobState {
 
     fn write_raw_persistent_data(
         drv: &mut FlashDriver,
-        buf: &mut ApobBufs,
+        buf: &mut HfBufs,
         data: ApobRawPersistentData,
         meta: Meta,
     ) {
         let mut found: Option<FlashAddr> = None;
         for offset in (0..APOB_META_SIZE).step_by(APOB_PERSISTENT_DATA_STRIDE) {
             let addr = meta.flash_addr(offset).unwrap_lite();
             // Infallible when using a slice
-            drv.flash_read(addr, &mut buf.persistent_data.as_mut_slice())
+            drv.flash_read(addr, &mut buf.apob_persistent_data.as_mut_slice())
                 .unwrap_lite();
-            if buf.persistent_data.iter().all(|c| *c == 0xFF) {
+            if buf.apob_persistent_data.iter().all(|c| *c == 0xFF) {
                 found = Some(addr);
                 break;
             }
@@ -873,6 +846,6 @@ impl ApobState {
             addr
         });
         // Infallible when using a slice
-        drv.flash_write(addr, &mut data.as_bytes()).unwrap_lite();
+        drv.flash_write(addr, data.as_bytes()).unwrap_lite();
     }
 }

drv/cosmo-hf/src/hf.rs

@@ -9,15 +9,16 @@ use drv_hf_api::{
     HF_PERSISTENT_DATA_STRIDE,
 };
 use idol_runtime::{
-    LeaseBufReader, LeaseBufWriter, Leased, LenLimit, NotificationHandler,
+    ClientError, LeaseBufWriter, Leased, LenLimit, NotificationHandler,
     RequestError, R, W,
 };
 use ringbuf::ringbuf_entry_root as ringbuf_entry;
 use userlib::{set_timer_relative, task_slot, RecvMessage, UnwrapLite};
 use zerocopy::{FromZeros, IntoBytes};
 
 use crate::{
-    apob, FlashAddr, FlashDriver, Trace, PAGE_SIZE_BYTES, SECTOR_SIZE_BYTES,
+    apob, apob::APOB_PERSISTENT_DATA_STRIDE, FlashAddr, FlashDriver, Trace,
+    PAGE_SIZE_BYTES, SECTOR_SIZE_BYTES,
 };
 
 task_slot!(HASH, hash_driver);
@@ -33,7 +34,36 @@ pub struct ServerImpl {
     hash: HashData,
 
     pub(crate) apob_state: apob::ApobState,
-    pub(crate) apob_buf: apob::ApobBufs,
+    pub(crate) buf: HfBufs,
+}
+
+pub(crate) struct HfBufs {
+    pub(crate) apob_persistent_data:
+        &'static mut [u8; APOB_PERSISTENT_DATA_STRIDE],
+    pub(crate) page: &'static mut [u8; PAGE_SIZE_BYTES],
+    pub(crate) scratch: &'static mut [u8; PAGE_SIZE_BYTES],
+}
+
+/// Grabs references to the static buffers.  Can only be called once.
+impl HfBufs {
+    pub fn claim_statics() -> Self {
+        use static_cell::ClaimOnceCell;
+        static BUFS: ClaimOnceCell<(
+            [u8; APOB_PERSISTENT_DATA_STRIDE],
+            [u8; PAGE_SIZE_BYTES],
+            [u8; PAGE_SIZE_BYTES],
+        )> = ClaimOnceCell::new((
+            [0; APOB_PERSISTENT_DATA_STRIDE],
+            [0; PAGE_SIZE_BYTES],
+            [0; PAGE_SIZE_BYTES],
+        ));
+        let (apob_persistent_data, page, scratch) = BUFS.claim();
+        Self {
+            apob_persistent_data,
+            page,
+            scratch,
+        }
+    }
 }
 
 /// This tunes how many bytes we hash in a single async timer notification
@@ -49,15 +79,15 @@ impl ServerImpl {
     /// Persistent data is loaded from the flash chip and used to select `dev`;
     /// in addition, it is made redundant (written to both virtual devices).
     pub fn new(mut drv: FlashDriver) -> Self {
-        let mut apob_buf = apob::ApobBufs::claim_statics();
-        let apob_state = apob::ApobState::init(&mut drv, &mut apob_buf);
+        let mut buf = HfBufs::claim_statics();
+        let apob_state = apob::ApobState::init(&mut drv, &mut buf);
 
         let mut out = Self {
             dev: drv_hf_api::HfDevSelect::Flash0,
             drv,
             hash: HashData::new(HASH.get_task_id()),
             apob_state,
-            apob_buf,
+            buf,
         };
         out.drv.set_flash_mux_state(HfMuxState::SP);
         out.ensure_persistent_data_is_redundant();
@@ -205,7 +235,7 @@ impl ServerImpl {
         };
         // flash_write is infallible when given a slice
         self.drv
-            .flash_write(addr, &mut raw_data.as_bytes())
+            .flash_write(addr, raw_data.as_bytes())
             .unwrap_lite();
     }
 
@@ -421,11 +451,12 @@ impl idl::InOrderHostFlashImpl for ServerImpl {
         self.drv.check_flash_mux_state()?;
         let addr = self.flash_addr(addr, data.len() as u32)?;
         self.invalidate_write();
+        // Read the entire data block into our address space.
+        data.read_range(0..data.len(), &mut self.buf.scratch[..data.len()])
+            .map_err(|_| RequestError::Fail(ClientError::WentAway))?;
+
         self.drv
-            .flash_write(
-                addr,
-                &mut LeaseBufReader::<_, 32>::from(data.into_inner()),
-            )
+            .flash_write(addr, &self.buf.scratch[..data.len()])
             .map_err(|()| RequestError::went_away())
     }
 
@@ -527,7 +558,7 @@ impl idl::InOrderHostFlashImpl for ServerImpl {
         data: Leased<R, [u8]>,
     ) -> Result<(), RequestError<drv_hf_api::ApobWriteError>> {
         self.apob_state
-            .write(&mut self.drv, &mut self.apob_buf, offset, data)
+            .write(&mut self.drv, &mut self.buf, offset, data)
             .map_err(RequestError::from)
     }
 
@@ -536,7 +567,7 @@ impl idl::InOrderHostFlashImpl for ServerImpl {
         _: &RecvMessage,
     ) -> Result<(), RequestError<drv_hf_api::ApobCommitError>> {
         self.apob_state
-            .commit(&mut self.drv, &mut self.apob_buf)
+            .commit(&mut self.drv, &mut self.buf)
             .map_err(RequestError::from)
     }
 
@@ -555,7 +586,7 @@ impl idl::InOrderHostFlashImpl for ServerImpl {
         data: Leased<W, [u8]>,
     ) -> Result<usize, RequestError<drv_hf_api::ApobReadError>> {
         self.apob_state
-            .read(&mut self.drv, &mut self.apob_buf, offset, data)
+            .read(&mut self.drv, &mut self.buf, offset, data)
             .map_err(RequestError::from)
     }
 
@@ -589,7 +620,7 @@ impl idl::InOrderHostFlashImpl for ServerImpl {
             // This also unlocks the APOB so it can be written (once muxed back
             // to the SP).
             self.apob_state =
-                apob::ApobState::init(&mut self.drv, &mut self.apob_buf);
+                apob::ApobState::init(&mut self.drv, &mut self.buf);
         }
         self.drv.set_flash_mux_state(state);
         self.invalidate_mux_switch();

drv/cosmo-hf/src/main.rs

@@ -345,44 +345,38 @@ impl FlashDriver {
         Ok(())
     }
 
-    /// Writes data from a `BufReader` into the flash
+    /// Writes data from a slice into the flash
     ///
-    /// This function will only return an error if it fails to write to a
-    /// provided lease; when given a slice, it is infallible.
-    fn flash_write(
-        &mut self,
-        addr: FlashAddr,
-        data: &mut dyn idol_runtime::BufReader<'_>,
-    ) -> Result<(), ()> {
-        loop {
-            let len = data.remaining_size().min(PAGE_SIZE_BYTES);
-            if len == 0 {
-                break;
-            }
-            self.flash_write_enable();
-            self.drv.data_bytes.set_count(len as u16);
-            self.drv.addr.set_addr(addr.0);
-            self.drv.dummy_cycles.set_count(0);
-            for i in 0..len.div_ceil(4) {
-                let mut v = [0u8; 4];
-                for (j, byte) in v.iter_mut().enumerate() {
-                    let k = i * 4 + j;
-                    if k < len {
-                        let Some(d) = data.read() else {
-                            return Err(());
-                        };
-                        *byte = d;
-                    }
-                }
-                let v = u32::from_le_bytes(v);
-                self.drv.tx_fifo_wdata.set_fifo_data(v);
-            }
-            self.drv.instr.set_opcode(instr::QUAD_INPUT_PAGE_PROGRAM_4B);
-            self.wait_fpga_busy();
+    /// This function will only return an error if the slice len is greater than
+    /// a page
+    fn flash_write(&mut self, addr: FlashAddr, data: &[u8]) -> Result<(), ()> {
+        // Callers are already doing everything at a page granularity
+        if data.len() > PAGE_SIZE_BYTES {
+            return Err(());
+        }
 
-            // Wait for the busy flag to be unset
-            self.wait_flash_busy(Trace::WriteBusy);
+        let len = data.len();
+
+        // Don't bother writing erased pages
+        if data.iter().all(|x| *x == 0xff) {
+            return Ok(());
         }
+
+        self.flash_write_enable();
+        self.drv.data_bytes.set_count(len as u16);
+        self.drv.addr.set_addr(addr.0);
+        self.drv.dummy_cycles.set_count(0);
+
+        for chunk in data.chunks(4) {
+            let v = u32::from_le_bytes(chunk.try_into().unwrap_lite());
+            self.drv.tx_fifo_wdata.set_fifo_data(v);
+        }
+
+        self.drv.instr.set_opcode(instr::QUAD_INPUT_PAGE_PROGRAM_4B);
+        self.wait_fpga_busy();
+
+        // Wait for the busy flag to be unset
+        self.wait_flash_busy(Trace::WriteBusy);
         Ok(())
     }