drtio_proto: add allocate step for flashing

This avoids reallocation when transfering binaries. Reallocation during flash handshakes could cause DRTIO timeouts.

artiq/firmware/libproto_artiq/drtioaux_proto.rs

@@ -139,7 +139,8 @@ pub enum Packet {
     CoreMgmtConfigEraseRequest { destination: u8 },
     CoreMgmtRebootRequest { destination: u8 },
     CoreMgmtAllocatorDebugRequest { destination: u8 },
-    CoreMgmtFlashRequest { destination: u8, last: bool, length: u16, data: [u8; MASTER_PAYLOAD_MAX_SIZE] },
+    CoreMgmtFlashRequest { destination: u8, payload_length: u32 },
+    CoreMgmtFlashAddDataRequest { destination: u8, last: bool, length: u16, data: [u8; MASTER_PAYLOAD_MAX_SIZE] },
     CoreMgmtDropLinkAck { destination: u8 },
     CoreMgmtDropLink,
     CoreMgmtGetLogReply { last: bool, length: u16, data: [u8; SAT_PAYLOAD_MAX_SIZE] },
@@ -485,24 +486,28 @@ impl Packet {
             0xda => Packet::CoreMgmtAllocatorDebugRequest {
                 destination: reader.read_u8()?,
             },
-            0xdb => {
+            0xdb => Packet::CoreMgmtFlashRequest {
+                destination: reader.read_u8()?,
+                payload_length: reader.read_u32()?,
+            },
+            0xdc => {
                 let destination = reader.read_u8()?;
                 let last = reader.read_bool()?;
                 let length = reader.read_u16()?;
                 let mut data: [u8; MASTER_PAYLOAD_MAX_SIZE] = [0; MASTER_PAYLOAD_MAX_SIZE];
                 reader.read_exact(&mut data[0..length as usize])?;
-                Packet::CoreMgmtFlashRequest {
+                Packet::CoreMgmtFlashAddDataRequest {
                     destination: destination,
                     last: last,
                     length: length,
                     data: data,
                 }
             },
-            0xdc => Packet::CoreMgmtDropLinkAck {
+            0xdd => Packet::CoreMgmtDropLinkAck {
                 destination: reader.read_u8()?,
             },
-            0xdd => Packet::CoreMgmtDropLink,
-            0xde => {
+            0xde => Packet::CoreMgmtDropLink,
+            0xdf => {
                 let last = reader.read_bool()?;
                 let length = reader.read_u16()?;
                 let mut data: [u8; SAT_PAYLOAD_MAX_SIZE] = [0; SAT_PAYLOAD_MAX_SIZE];
@@ -513,7 +518,7 @@ impl Packet {
                     data: data,
                 }
             },
-            0xdf => {
+            0xe0 => {
                 let last = reader.read_bool()?;
                 let length = reader.read_u16()?;
                 let mut value: [u8; SAT_PAYLOAD_MAX_SIZE] = [0; SAT_PAYLOAD_MAX_SIZE];
@@ -524,7 +529,7 @@ impl Packet {
                     value: value,
                 }
             },
-            0xe0 => Packet::CoreMgmtReply {
+            0xe1 => Packet::CoreMgmtReply {
                 succeeded: reader.read_bool()?,
             },
 
@@ -884,33 +889,38 @@ impl Packet {
                 writer.write_u8(0xda)?;
                 writer.write_u8(destination)?;
             },
-            Packet::CoreMgmtFlashRequest { destination, last, length, data } => {
+            Packet::CoreMgmtFlashRequest { destination, payload_length } => {
                 writer.write_u8(0xdb)?;
                 writer.write_u8(destination)?;
+                writer.write_u32(payload_length)?;
+            },
+            Packet::CoreMgmtFlashAddDataRequest { destination, last, length, data } => {
+                writer.write_u8(0xdc)?;
+                writer.write_u8(destination)?;
                 writer.write_bool(last)?;
                 writer.write_u16(length)?;
                 writer.write_all(&data[..length as usize])?;
             },
             Packet::CoreMgmtDropLinkAck { destination } => {
-                writer.write_u8(0xdc)?;
+                writer.write_u8(0xdd)?;
                 writer.write_u8(destination)?;
             },
             Packet::CoreMgmtDropLink => 
-                writer.write_u8(0xdd)?,
+                writer.write_u8(0xde)?,
             Packet::CoreMgmtGetLogReply { last, length, data } => {
-                writer.write_u8(0xde)?;
+                writer.write_u8(0xdf)?;
                 writer.write_bool(last)?;
                 writer.write_u16(length)?;
                 writer.write_all(&data[0..length as usize])?;
             },
             Packet::CoreMgmtConfigReadReply { last, length, value } => {
-                writer.write_u8(0xdf)?;
+                writer.write_u8(0xe0)?;
                 writer.write_bool(last)?;
                 writer.write_u16(length)?;
                 writer.write_all(&value[0..length as usize])?;
             },
             Packet::CoreMgmtReply { succeeded } => {
-                writer.write_u8(0xe0)?;
+                writer.write_u8(0xe1)?;
                 writer.write_bool(succeeded)?;
             },
         }

artiq/firmware/runtime/mgmt.rs

@@ -546,9 +546,29 @@ mod remote_coremgmt {
         routing_table: &RoutingTable, linkno: u8,
         destination: u8, stream: &mut TcpStream, image: &[u8]) -> Result<(), Error<SchedError>> {
 
+        let alloc_reply = drtio::aux_transact(io, aux_mutex, ddma_mutex, subkernel_mutex, routing_table, linkno,
+            &Packet::CoreMgmtFlashRequest {
+                destination: destination,
+                payload_length: image.len() as u32,
+            });
+
+        match alloc_reply {
+            Ok(Packet::CoreMgmtReply { succeeded: true }) => Ok(()),
+            Ok(packet) => {
+                error!("received unexpected aux packet: {:?}", packet);
+                Reply::Error.write_to(stream)?;
+                Err(drtio::Error::UnexpectedReply)
+            }
+            Err(e) => {
+                error!("aux packet error ({})", e);
+                Reply::Error.write_to(stream)?;
+                Err(e)
+            }
+        }?;
+
         match drtio::partition_data(&image, |slice, status, len: usize| {
             let reply = drtio::aux_transact(io, aux_mutex, ddma_mutex, subkernel_mutex, routing_table, linkno, 
-                &Packet::CoreMgmtFlashRequest {
+                &Packet::CoreMgmtFlashAddDataRequest {
                     destination: destination, length: len as u16, last: status.is_last(), data: *slice});
             match reply {
                 Ok(Packet::CoreMgmtReply { succeeded: true }) => Ok(()),

artiq/firmware/satman/main.rs

@@ -628,7 +628,13 @@ fn process_aux_packet(dmamgr: &mut DmaManager, analyzer: &mut Analyzer, kernelmg
             warn!("restarting");
             unsafe { spiflash::reload(); }
         }
-        drtioaux::Packet::CoreMgmtFlashRequest { destination: _destination, last, length, data } => {
+        drtioaux::Packet::CoreMgmtFlashRequest { destination: _destination, payload_length } => {
+            forward!(router, _routing_table, _destination, *rank, *self_destination, _repeaters, &packet);
+
+            coremgr.allocate_image_buffer(payload_length as usize);
+            drtioaux::send(0, &drtioaux::Packet::CoreMgmtReply { succeeded: true })
+        }
+        drtioaux::Packet::CoreMgmtFlashAddDataRequest { destination: _destination, last, length, data } => {
             forward!(router, _routing_table, _destination, *rank, *self_destination, _repeaters, &packet);
 
             coremgr.add_image_data(&data, length as usize);

artiq/firmware/satman/mgmt.rs

@@ -106,6 +106,10 @@ impl Manager {
         drtioaux::send(0, &drtioaux::Packet::CoreMgmtReply { succeeded })
     }
 
+    pub fn allocate_image_buffer(&mut self, image_size: usize) {
+        self.image_payload = Cursor::new(Vec::with_capacity(image_size));
+    }
+
     pub fn add_image_data(&mut self, data: &[u8], data_len: usize) {
         self.image_payload.write_all(&data[..data_len]).unwrap();
     }