From 3c6b2ef2e2c232941674fc269ee467a65195d3d2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marc-Andr=C3=A9=20Lureau?= <marcandre.lureau@redhat.com>
Date: Wed, 28 Aug 2024 15:34:55 +0400
Subject: [PATCH] feat(server): move the encoder to blocking task
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Since it is a CPU intesive task, we should spawn a task to avoid
blocking other async tasks.

The UpdateFragmenterOwned is quite a gross hack to allow returning the result
from a the task lifetime. I don't know how to accomplish this better, or
else we have to add some Arc<Mutex> stuff.

Signed-off-by: Marc-André Lureau <marcandre.lureau@redhat.com>
---
 crates/ironrdp-server/src/encoder/mod.rs | 22 ++++++++++++++++++++++
 crates/ironrdp-server/src/server.rs      | 20 ++++++++++++++------
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/crates/ironrdp-server/src/encoder/mod.rs b/crates/ironrdp-server/src/encoder/mod.rs
index 0bd44e7d..0ee99272 100644
--- a/crates/ironrdp-server/src/encoder/mod.rs
+++ b/crates/ironrdp-server/src/encoder/mod.rs
@@ -131,6 +131,14 @@ impl UpdateEncoder {
         update(self, bitmap)
     }
 
+    pub(crate) fn fragmenter_from_owned(&self, res: UpdateFragmenterOwned) -> UpdateFragmenter<'_> {
+        UpdateFragmenter {
+            code: res.code,
+            index: res.index,
+            data: &self.buffer[0..res.len],
+        }
+    }
+
     fn bitmap_update(&mut self, bitmap: BitmapUpdate) -> Result<UpdateFragmenter<'_>> {
         let len = loop {
             match self.bitmap.encode(&bitmap, self.buffer.as_mut_slice()) {
@@ -208,6 +216,12 @@ impl UpdateEncoder {
     }
 }
 
+pub(crate) struct UpdateFragmenterOwned {
+    code: UpdateCode,
+    index: usize,
+    len: usize,
+}
+
 pub(crate) struct UpdateFragmenter<'a> {
     code: UpdateCode,
     index: usize,
@@ -219,6 +233,14 @@ impl<'a> UpdateFragmenter<'a> {
         Self { code, index: 0, data }
     }
 
+    pub(crate) fn into_owned(self) -> UpdateFragmenterOwned {
+        UpdateFragmenterOwned {
+            code: self.code,
+            index: self.index,
+            len: self.data.len(),
+        }
+    }
+
     pub(crate) fn size_hint(&self) -> usize {
         FASTPATH_HEADER_SIZE + cmp::min(self.data.len(), MAX_FASTPATH_UPDATE_SIZE)
     }
diff --git a/crates/ironrdp-server/src/server.rs b/crates/ironrdp-server/src/server.rs
index c8fc2506..752b63eb 100644
--- a/crates/ironrdp-server/src/server.rs
+++ b/crates/ironrdp-server/src/server.rs
@@ -376,13 +376,21 @@ impl RdpServer {
         user_channel_id: u16,
         io_channel_id: u16,
         buffer: &mut Vec<u8>,
-        encoder: &mut UpdateEncoder,
-    ) -> Result<RunState>
+        mut encoder: UpdateEncoder,
+    ) -> Result<(RunState, UpdateEncoder)>
     where
         W: FramedWrite,
     {
         let mut fragmenter = match update {
-            DisplayUpdate::Bitmap(bitmap) => encoder.bitmap(bitmap),
+            DisplayUpdate::Bitmap(bitmap) => {
+                let (enc, res) = task::spawn_blocking(move || {
+                    let res = encoder.bitmap(bitmap).map(|r| r.into_owned());
+                    (encoder, res)
+                })
+                .await?;
+                encoder = enc;
+                res.map(|r| encoder.fragmenter_from_owned(r))
+            }
             DisplayUpdate::PointerPosition(pos) => encoder.pointer_position(pos),
             DisplayUpdate::Resize(desktop_size) => {
                 debug!(?desktop_size, "Display resize");
@@ -400,7 +408,7 @@ impl RdpServer {
                 };
                 let msg = encode_vec(&X224(pdu))?;
                 writer.write_all(&msg).await?;
-                return Ok(RunState::DeactivationReactivation { desktop_size });
+                return Ok((RunState::DeactivationReactivation { desktop_size }, encoder));
             }
             DisplayUpdate::RGBAPointer(ptr) => encoder.rgba_pointer(ptr),
             DisplayUpdate::ColorPointer(ptr) => encoder.color_pointer(ptr),
@@ -420,7 +428,7 @@ impl RdpServer {
                 .context("failed to write display update")?;
         }
 
-        Ok(RunState::Continue)
+        Ok((RunState::Continue, encoder))
     }
 
     async fn dispatch_server_events<W>(
@@ -526,7 +534,7 @@ impl RdpServer {
                 },
 
                 Some(update) = display_updates.next_update() => {
-                    state = self.dispatch_display_update(update, writer, user_channel_id, io_channel_id, &mut buffer, &mut encoder).await?;
+                    (state, encoder) = self.dispatch_display_update(update, writer, user_channel_id, io_channel_id, &mut buffer, encoder).await?;
                 }
 
                 nevents = self.ev_receiver.recv_many(&mut events, 100) => {