Merge pull request #12 from noctisynth/perf/mem-and-inline

perf: add `#[inline]` and use chunks
noctisynth · May 17, 2024 · a5fae44 · a5fae44
2 parents 03c2f32 + 2729bf4
commit a5fae44
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 14 deletions.
diff --git a/src/models/packet.rs b/src/models/packet.rs
@@ -295,16 +295,16 @@ impl<'a> OED<'a> {
         stream.send(&self.plain_data()?).await?;
 
         self.chunk_count = 0;
-        let mut remaining_data = &self.encrypted_data[..];
-        while !remaining_data.is_empty() {
-            let chunk_size = remaining_data.len().min(1024);
-
-            let chunk_length = chunk_size as u32;
-
-            stream.send(&chunk_length.to_be_bytes()).await?;
-            stream.send(&remaining_data[..chunk_size]).await?;
-
-            remaining_data = &remaining_data[chunk_size..];
+        let mut chunks = self.encrypted_data.chunks(1024);
+        loop {
+            if let Some(chunk) = chunks.next() {
+                let chunk_size = chunk.len() as u32;
+                stream.send(&chunk_size.to_be_bytes()).await?;
+                stream.send(&chunk).await?;
+                self.chunk_count += 1;
+            } else {
+                break;
+            }
         }
         stream.send(&STOP_FLAG).await?;
 
@@ -320,7 +320,11 @@ impl<'a> OED<'a> {
         Ok(plain_bytes)
     }
 
-    pub fn get_data(&self) -> Vec<u8> {
-        self.data.clone().unwrap()
+    pub fn take(&mut self) -> Vec<u8> {
+        self.data.take().unwrap()
+    }
+
+    pub fn get_data(&self) -> &[u8] {
+        self.data.as_ref().unwrap()
     }
 }
diff --git a/src/models/session.rs b/src/models/session.rs
@@ -193,7 +193,7 @@ impl Session {
         let content = OED::new(&self.aes_key)
             .from_stream(socket)
             .await?
-            .get_data();
+            .take();
         let status_code = OSC::from_stream(socket).await?.status_code;
         let response = Response::new(None, content, None, status_code, flag);
 

diff --git a/src/utils/gear.rs b/src/utils/gear.rs
@@ -50,16 +50,17 @@ impl Socket {
         }
     }
 
+    #[inline]
     pub async fn peer_addr(&self) -> Result<SocketAddr> {
         Ok(self.writer.lock().await.peer_addr()?)
     }
 
+    #[inline]
     pub async fn recv_usize(&self) -> Result<usize> {
         let mut len_bytes = [0; 4];
         #[cfg(not(feature = "perf"))]
         self.reader.lock().await.read_exact(&mut len_bytes).await?;
         #[cfg(feature = "perf")]
-        #[cfg(feature = "perf")]
         {
             use colored::Colorize;
             let now = tokio::time::Instant::now();
@@ -73,24 +74,28 @@ impl Socket {
         Ok(u32::from_be_bytes(len_bytes) as usize)
     }
 
+    #[inline]
     pub async fn recv_u32(&self) -> Result<u32> {
         let mut len_bytes = [0; 4];
         self.reader.lock().await.read_exact(&mut len_bytes).await?;
         Ok(u32::from_be_bytes(len_bytes))
     }
 
+    #[inline]
     pub async fn recv(&self, len: usize) -> Result<Vec<u8>> {
         let mut recv_bytes: Vec<u8> = vec![0; len];
         self.reader.lock().await.read_exact(&mut recv_bytes).await?;
         Ok(recv_bytes)
     }
 
+    #[inline]
     pub async fn recv_str(&self, len: usize) -> Result<String> {
         let mut recv_bytes: Vec<u8> = vec![0; len];
         self.reader.lock().await.read_exact(&mut recv_bytes).await?;
         Ok(String::from_utf8(recv_bytes)?)
     }
 
+    #[inline]
     pub async fn send(&self, data: &[u8]) -> Result<()> {
         let mut writer = self.writer.lock().await;
         writer.write_all(&data).await?;