From f6cf66966ddd28e40f56bfb95212cb536b838134 Mon Sep 17 00:00:00 2001
From: occheung <dc@m-labs.hk>
Date: Thu, 29 Aug 2024 12:30:15 +0800
Subject: [PATCH] remote coremgmt: restart device gracefully after flashing

---
 .../firmware/libproto_artiq/drtioaux_proto.rs |  3 ++
 artiq/firmware/runtime/mgmt.rs                |  1 +
 artiq/firmware/satman/main.rs                 | 11 +---
 artiq/firmware/satman/mgmt.rs                 | 50 +++++++++++++++----
 4 files changed, 46 insertions(+), 19 deletions(-)

diff --git a/artiq/firmware/libproto_artiq/drtioaux_proto.rs b/artiq/firmware/libproto_artiq/drtioaux_proto.rs
index ba826873a..8dbe9c98a 100644
--- a/artiq/firmware/libproto_artiq/drtioaux_proto.rs
+++ b/artiq/firmware/libproto_artiq/drtioaux_proto.rs
@@ -143,6 +143,7 @@ pub enum Packet {
     CoreMgmtConfigReadReply { last: bool, length: u16, value: [u8; SAT_PAYLOAD_MAX_SIZE] },
     CoreMgmtAck,
     CoreMgmtNack,
+    CoreMgmtRebootImminent,
 }
 
 impl Packet {
@@ -507,6 +508,7 @@ impl Packet {
             },
             0xdd => Packet::CoreMgmtAck,
             0xde => Packet::CoreMgmtNack,
+            0xdf => Packet::CoreMgmtRebootImminent,
 
             ty => return Err(Error::UnknownPacket(ty))
         })
@@ -878,6 +880,7 @@ impl Packet {
             },
             Packet::CoreMgmtAck => writer.write_u8(0xdd)?,
             Packet::CoreMgmtNack => writer.write_u8(0xde)?,
+            Packet::CoreMgmtRebootImminent => writer.write_u8(0xdf)?,
         }
         Ok(())
     }
diff --git a/artiq/firmware/runtime/mgmt.rs b/artiq/firmware/runtime/mgmt.rs
index 95ac3870b..34a2c7d2e 100644
--- a/artiq/firmware/runtime/mgmt.rs
+++ b/artiq/firmware/runtime/mgmt.rs
@@ -372,6 +372,7 @@ mod remote_coremgmt {
                     destination: destination, length: len as u16, last: status.is_last(), data: *slice});
             match reply {
                 Ok(Packet::CoreMgmtAck) => Ok(()),
+                Ok(Packet::CoreMgmtRebootImminent) if status.is_last() => Ok(()),
                 Ok(packet) => {
                     error!("received unexpected aux packet: {:?}", packet);
                     Err(drtio::Error::UnexpectedReply)
diff --git a/artiq/firmware/satman/main.rs b/artiq/firmware/satman/main.rs
index 9513f1a81..c7ca63775 100644
--- a/artiq/firmware/satman/main.rs
+++ b/artiq/firmware/satman/main.rs
@@ -562,17 +562,10 @@ fn process_aux_packet(dmamgr: &mut DmaManager, analyzer: &mut Analyzer, kernelmg
             forward!(router, _routing_table, _destination, *rank, *self_destination, _repeaters, &packet);
 
             coremgr.add_data(&data, length as usize);
-            let mut succeeded = true;
             if last {
-                succeeded = coremgr.write_config().is_ok();
-                debug!("Write succeeded: {}", succeeded);
-                coremgr.clear_data();
-            }
-
-            if succeeded {
-                drtioaux::send(0, &drtioaux::Packet::CoreMgmtAck)
+                coremgr.write_config()
             } else {
-                drtioaux::send(0, &drtioaux::Packet::CoreMgmtNack)
+                drtioaux::send(0, &drtioaux::Packet::CoreMgmtAck)
             }
         }
         drtioaux::Packet::CoreMgmtConfigRemoveRequest { destination: _destination, length, key } => {
diff --git a/artiq/firmware/satman/mgmt.rs b/artiq/firmware/satman/mgmt.rs
index 245e5f4d6..7e220fcd9 100644
--- a/artiq/firmware/satman/mgmt.rs
+++ b/artiq/firmware/satman/mgmt.rs
@@ -1,13 +1,12 @@
 use alloc::vec::Vec;
 
 use routing::{Sliceable, SliceMeta};
-use board_misoc::config;
+use board_artiq::drtioaux;
+use board_misoc::{clock, config, csr, spiflash};
 use io::{Cursor, ProtoRead, ProtoWrite};
 use proto_artiq::drtioaux_proto::SAT_PAYLOAD_MAX_SIZE;
 
 
-type Result<T> = core::result::Result<T, ()>;
-
 pub struct Manager {
     current_payload: Cursor<Vec<u8>>,
     last_value: Sliceable,
@@ -21,7 +20,7 @@ impl Manager {
         }
     }
 
-    pub fn fetch_config_value(&mut self, key: &str) -> Result<()> {
+    pub fn fetch_config_value(&mut self, key: &str) -> Result<(), ()> {
         config::read(key, |result| result.map(
             |value| self.last_value = Sliceable::new(0, value.to_vec())
         )).map_err(|_err| warn!("read error: no such key"))
@@ -40,13 +39,44 @@ impl Manager {
         self.current_payload.set_position(0);
     }
 
-    pub fn write_config(&mut self) -> Result<()> {
-        let key = self.current_payload.read_string().map_err(
-            |err| error!("error on reading key: {:?}", err))?;
+    pub fn write_config(&mut self) -> Result<(), drtioaux::Error<!>> {
+        let key = match self.current_payload.read_string() {
+            Ok(key) => key,
+            Err(err) => {
+                self.clear_data();
+                error!("error on reading key: {:?}", err);
+                return drtioaux::send(0, &drtioaux::Packet::CoreMgmtNack);
+            }
+        };
+
         let value = self.current_payload.read_bytes().unwrap();
 
-        config::write(&key, &value).map_err(|err| {
-            error!("error on writing config: {:?}", err);
-        })
+        match key.as_str() {
+            "gateware" | "bootloader" | "firmware" => {
+                drtioaux::send(0, &drtioaux::Packet::CoreMgmtRebootImminent)?;
+                #[cfg(not(soc_platform = "efc"))]
+                unsafe {
+                    clock::spin_us(10000);
+                    csr::gt_drtio::txenable_write(0);
+                }
+                config::write(&key, &value).expect("failed to write to flash storage");
+                warn!("restarting");
+                unsafe { spiflash::reload(); }
+            }
+
+            _ => {
+                let succeeded = config::write(&key, &value).map_err(|err| {
+                    error!("error on writing config: {:?}", err);
+                }).is_ok();
+
+                self.clear_data();
+
+                if succeeded {
+                    drtioaux::send(0, &drtioaux::Packet::CoreMgmtAck)
+                } else {
+                    drtioaux::send(0, &drtioaux::Packet::CoreMgmtNack)
+                }
+            }
+        }
     }
 }