Use thingbuf instead of a custom "m"psc.

Author: veluca93

pixie-shared/src/lib.rs

@@ -39,7 +39,7 @@ pub type ChunkHash = [u8; OUT_LEN];
 pub type Offset = usize;
 
 /// Describes one chunk from a disk.
-#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
 pub struct Chunk {
     pub hash: ChunkHash,
     pub start: Offset,

pixie-uefi/Cargo.lock

@@ -27,6 +27,7 @@ postcard = { version = "1.1.3", default-features = false, features = ["alloc"] }
 rand = { version = "0.8.5", default-features = false }
 rand_xoshiro = { version = "0.6.0", default-features = false }
 smoltcp = { version = "0.12.0", default-features = false, features = ["alloc", "proto-ipv4", "medium-ethernet", "socket-udp", "socket-tcp", "socket-dhcpv4", "async", "socket-tcp-cubic"] }
+thingbuf = { version = "0.1.6", default-features = false, features = ["alloc"] }
 uefi = { version = "0.33.0", features = ["alloc", "global_allocator", "panic_handler"] }
 
 [dependencies.pixie-shared]

pixie-uefi/Cargo.toml

@@ -1,7 +1,7 @@
 use crate::{
     os::{
         error::{Error, Result},
-        mpsc, TcpStream, UefiOS, PACKET_SIZE,
+        TcpStream, UefiOS, PACKET_SIZE,
     },
     MIN_MEMORY,
 };
@@ -162,10 +162,10 @@ pub async fn flash(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> {
 
     let mut received = BTreeMap::new();
 
-    let (tx, mut rx) = mpsc::channel(128);
+    let (tx, rx) = thingbuf::mpsc::channel(128);
 
     let task1 = async {
-        let mut tx = tx;
+        let tx = tx;
         let mut last_seen = Vec::new();
         let total_mem = os.get_total_mem();
         let max_chunks = (total_mem.saturating_sub(MIN_MEMORY) as usize / MAX_CHUNK_SIZE).max(128);
@@ -184,7 +184,7 @@ pub async fn flash(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> {
                     let chunk =
                         handle_packet(buf, &mut chunks_info, &mut received, &mut last_seen)?;
                     if let Some((pos, data)) = chunk {
-                        tx.send((pos, data)).await;
+                        tx.send((pos, data)).await.expect("receiver was dropped");
                     }
 
                     assert_eq!(last_seen.len(), received.len());

pixie-uefi/src/flash.rs

@@ -48,7 +48,6 @@ mod boot_options;
 pub mod disk;
 pub mod error;
 mod executor;
-pub mod mpsc;
 mod net;
 mod rng;
 mod sync;

pixie-uefi/src/os/mod.rs

@@ -1,7 +1,7 @@
 use crate::{
     os::{
         error::{Error, Result},
-        mpsc, TcpStream, UefiOS,
+        TcpStream, UefiOS,
     },
     parse_disk, MIN_MEMORY,
 };
@@ -92,13 +92,13 @@ pub async fn store(os: UefiOS, server_address: SocketAddrV4) -> Result<()> {
         BytesFmt(total_mem)
     );
 
-    let (tx1, mut rx1) = mpsc::channel(channel_size);
-    let (tx2, mut rx2) = mpsc::channel(channel_size);
-    let (tx3, mut rx3) = mpsc::channel(channel_size);
-    let (tx4, mut rx4) = mpsc::channel(channel_size);
+    let (tx1, rx1) = thingbuf::mpsc::channel(channel_size);
+    let (tx2, rx2) = thingbuf::mpsc::channel(channel_size);
+    let (tx3, rx3) = thingbuf::mpsc::channel(channel_size);
+    let (tx4, rx4) = thingbuf::mpsc::channel(channel_size);
 
     let task1 = async {
-        let mut tx1 = tx1;
+        let tx1 = tx1;
         for chunk_info in chunks {
             let mut data = vec![0; chunk_info.size];
             disk.read(chunk_info.start as u64, &mut data).await?;
@@ -110,45 +110,49 @@ pub async fn store(os: UefiOS, server_address: SocketAddrV4) -> Result<()> {
                 size: chunk_info.size,
                 csize: cdata.len(),
             };
-            tx1.send((chunk, cdata)).await;
+            tx1.send((chunk, cdata)).await.expect("receiver dropped");
         }
         Ok::<_, Error>(())
     };
 
     let task2 = async {
-        let mut tx2 = tx2;
+        let tx2 = tx2;
         while let Some((chunk, cdata)) = rx1.recv().await {
             let req = TcpRequest::HasChunk(chunk.hash);
             let buf = postcard::to_allocvec(&req)?;
             stream_get_csize.send_u64_le(buf.len() as u64).await?;
             stream_get_csize.send(&buf).await?;
-            tx2.send((chunk, cdata)).await;
+            tx2.send((chunk, cdata)).await.expect("receiver dropped");
         }
         Ok(())
     };
 
     let task3 = async {
-        let mut tx3 = tx3;
+        let tx3 = tx3;
         while let Some((chunk, cdata)) = rx2.recv().await {
             let len = stream_get_csize.recv_u64_le().await?;
             let mut buf = vec![0; len as usize];
             stream_get_csize.recv(&mut buf).await?;
             let has_chunk: bool = postcard::from_bytes(&buf)?;
-            tx3.send((chunk, cdata, has_chunk)).await;
+            tx3.send((chunk, cdata, has_chunk))
+                .await
+                .expect("receiver dropped");
         }
         Ok(())
     };
 
     let task4 = async {
-        let mut tx4 = tx4;
+        let tx4 = tx4;
         while let Some((chunk, cdata, has_chunk)) = rx3.recv().await {
             if !has_chunk {
                 let req = TcpRequest::UploadChunk(cdata);
                 let buf = postcard::to_allocvec(&req)?;
                 stream_upload_chunk.send_u64_le(buf.len() as u64).await?;
                 stream_upload_chunk.send(&buf).await?;
             }
-            tx4.send((chunk, has_chunk)).await;
+            tx4.send((chunk, has_chunk))
+                .await
+                .expect("receiver dropped");
         }
         Ok(())
     };