refactor(server): make UpdateEncoder::update() an iterator

A single display update can now result in multiple commands / update
code (FastPathUpdate).

The update dispatching and bitmap encoding is now done by the
UpdateEncoder itself.

Signed-off-by: Marc-André Lureau <marcandre.lureau@redhat.com>

Author: elmarco

crates/ironrdp-server/src/encoder/bitmap.rs

@@ -7,6 +7,7 @@ use ironrdp_pdu::geometry::InclusiveRectangle;
 use crate::BitmapUpdate;
 
 // PERF: we could also remove the need for this buffer
+#[derive(Clone)]
 pub(crate) struct BitmapEncoder {
     buffer: Vec<u8>,
 }

crates/ironrdp-server/src/encoder/mod.rs

@@ -12,7 +12,7 @@ use ironrdp_pdu::surface_commands::{ExtendedBitmapDataPdu, SurfaceBitsPdu, Surfa
 use self::bitmap::BitmapEncoder;
 use self::rfx::RfxEncoder;
 use super::BitmapUpdate;
-use crate::{ColorPointer, Framebuffer, RGBAPointer};
+use crate::{time_warn, ColorPointer, DisplayUpdate, Framebuffer, RGBAPointer};
 
 mod bitmap;
 mod fast_path;
@@ -59,12 +59,18 @@ impl UpdateEncoder {
         }
     }
 
+    pub(crate) fn update(&mut self, update: DisplayUpdate) -> EncoderIter<'_> {
+        EncoderIter {
+            encoder: self,
+            update: Some(update),
+        }
+    }
+
     pub(crate) fn set_desktop_size(&mut self, size: DesktopSize) {
         self.desktop_size = size;
     }
 
-    #[allow(clippy::unused_self)]
-    pub(crate) fn rgba_pointer(&mut self, ptr: RGBAPointer) -> Result<UpdateFragmenter> {
+    fn rgba_pointer(ptr: RGBAPointer) -> Result<UpdateFragmenter> {
         let xor_mask = ptr.data;
 
         let hot_spot = Point16 {
@@ -86,8 +92,7 @@ impl UpdateEncoder {
         Ok(UpdateFragmenter::new(UpdateCode::NewPointer, encode_vec(&ptr)?))
     }
 
-    #[allow(clippy::unused_self)]
-    pub(crate) fn color_pointer(&mut self, ptr: ColorPointer) -> Result<UpdateFragmenter> {
+    fn color_pointer(ptr: ColorPointer) -> Result<UpdateFragmenter> {
         let hot_spot = Point16 {
             x: ptr.hot_x,
             y: ptr.hot_y,
@@ -103,23 +108,26 @@ impl UpdateEncoder {
         Ok(UpdateFragmenter::new(UpdateCode::ColorPointer, encode_vec(&ptr)?))
     }
 
-    #[allow(clippy::unused_self)]
-    pub(crate) fn default_pointer(&mut self) -> Result<UpdateFragmenter> {
+    fn default_pointer() -> Result<UpdateFragmenter> {
         Ok(UpdateFragmenter::new(UpdateCode::DefaultPointer, vec![]))
     }
 
-    #[allow(clippy::unused_self)]
-    pub(crate) fn hide_pointer(&mut self) -> Result<UpdateFragmenter> {
+    fn hide_pointer() -> Result<UpdateFragmenter> {
         Ok(UpdateFragmenter::new(UpdateCode::HiddenPointer, vec![]))
     }
 
-    #[allow(clippy::unused_self)]
-    pub(crate) fn pointer_position(&mut self, pos: PointerPositionAttribute) -> Result<UpdateFragmenter> {
+    fn pointer_position(pos: PointerPositionAttribute) -> Result<UpdateFragmenter> {
         Ok(UpdateFragmenter::new(UpdateCode::PositionPointer, encode_vec(&pos)?))
     }
 
-    pub(crate) fn bitmap(&mut self, bitmap: BitmapUpdate) -> Result<UpdateFragmenter> {
-        let res = self.bitmap_updater.handle(&bitmap);
+    async fn bitmap(&mut self, bitmap: BitmapUpdate) -> Result<UpdateFragmenter> {
+        // Clone to satisfy spawn_blocking 'static requirement
+        // this should be cheap, even if using bitmap, since vec![] will be empty
+        let mut updater = self.bitmap_updater.clone();
+        let (res, bitmap) =
+            tokio::task::spawn_blocking(move || time_warn!("Encoding bitmap", 10, (updater.handle(&bitmap), bitmap)))
+                .await
+                .unwrap();
         if bitmap.x == 0
             && bitmap.y == 0
             && bitmap.width.get() == self.desktop_size.width
@@ -134,7 +142,31 @@ impl UpdateEncoder {
     }
 }
 
-#[derive(Debug)]
+pub(crate) struct EncoderIter<'a> {
+    encoder: &'a mut UpdateEncoder,
+    update: Option<DisplayUpdate>,
+}
+
+impl EncoderIter<'_> {
+    pub(crate) async fn next(&mut self) -> Option<Result<UpdateFragmenter>> {
+        let update = self.update.take()?;
+        let encoder = &mut self.encoder;
+
+        let res = match update {
+            DisplayUpdate::Bitmap(bitmap) => encoder.bitmap(bitmap).await,
+            DisplayUpdate::PointerPosition(pos) => UpdateEncoder::pointer_position(pos),
+            DisplayUpdate::RGBAPointer(ptr) => UpdateEncoder::rgba_pointer(ptr),
+            DisplayUpdate::ColorPointer(ptr) => UpdateEncoder::color_pointer(ptr),
+            DisplayUpdate::HidePointer => UpdateEncoder::hide_pointer(),
+            DisplayUpdate::DefaultPointer => UpdateEncoder::default_pointer(),
+            DisplayUpdate::Resize(_) => return None,
+        };
+
+        Some(res)
+    }
+}
+
+#[derive(Debug, Clone)]
 enum BitmapUpdater {
     None(NoneHandler),
     Bitmap(BitmapHandler),
@@ -155,7 +187,7 @@ trait BitmapUpdateHandler {
     fn handle(&mut self, bitmap: &BitmapUpdate) -> Result<UpdateFragmenter>;
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 struct NoneHandler;
 
 impl BitmapUpdateHandler for NoneHandler {
@@ -169,6 +201,7 @@ impl BitmapUpdateHandler for NoneHandler {
     }
 }
 
+#[derive(Clone)]
 struct BitmapHandler {
     bitmap: BitmapEncoder,
 }
@@ -209,7 +242,7 @@ impl BitmapUpdateHandler for BitmapHandler {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 struct RemoteFxHandler {
     remotefx: RfxEncoder,
     codec_id: u8,

crates/ironrdp-server/src/encoder/rfx.rs

@@ -11,7 +11,7 @@ use ironrdp_pdu::WriteCursor;
 
 use crate::BitmapUpdate;
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) struct RfxEncoder {
     entropy_algorithm: rfx::EntropyAlgorithm,
 }

crates/ironrdp-server/src/server.rs

@@ -32,7 +32,7 @@ use crate::clipboard::CliprdrServerFactory;
 use crate::display::{DisplayUpdate, RdpServerDisplay};
 use crate::encoder::UpdateEncoder;
 use crate::handler::RdpServerInputHandler;
-use crate::{builder, capabilities, time_warn, SoundServerFactory};
+use crate::{builder, capabilities, SoundServerFactory};
 
 #[derive(Clone)]
 pub struct RdpServerOptions {
@@ -423,39 +423,30 @@ impl RdpServer {
         buffer: &mut Vec<u8>,
         mut encoder: UpdateEncoder,
     ) -> Result<(RunState, UpdateEncoder)> {
-        let mut fragmenter = match update {
-            DisplayUpdate::Bitmap(bitmap) => {
-                let (enc, res) = task::spawn_blocking(move || {
-                    let res = time_warn!("Encoding bitmap", 10, encoder.bitmap(bitmap));
-                    (encoder, res)
-                })
-                .await?;
-                encoder = enc;
-                res
-            }
-            DisplayUpdate::PointerPosition(pos) => encoder.pointer_position(pos),
-            DisplayUpdate::Resize(desktop_size) => {
-                debug!(?desktop_size, "Display resize");
-                encoder.set_desktop_size(desktop_size);
-                deactivate_all(io_channel_id, user_channel_id, writer).await?;
-                return Ok((RunState::DeactivationReactivation { desktop_size }, encoder));
-            }
-            DisplayUpdate::RGBAPointer(ptr) => encoder.rgba_pointer(ptr),
-            DisplayUpdate::ColorPointer(ptr) => encoder.color_pointer(ptr),
-            DisplayUpdate::HidePointer => encoder.hide_pointer(),
-            DisplayUpdate::DefaultPointer => encoder.default_pointer(),
+        if let DisplayUpdate::Resize(desktop_size) = update {
+            debug!(?desktop_size, "Display resize");
+            encoder.set_desktop_size(desktop_size);
+            deactivate_all(io_channel_id, user_channel_id, writer).await?;
+            return Ok((RunState::DeactivationReactivation { desktop_size }, encoder));
         }
-        .context("error during update encoding")?;
 
-        if fragmenter.size_hint() > buffer.len() {
-            buffer.resize(fragmenter.size_hint(), 0);
-        }
+        let mut encoder_iter = encoder.update(update);
+        loop {
+            let Some(fragmenter) = encoder_iter.next().await else {
+                break;
+            };
 
-        while let Some(len) = fragmenter.next(buffer) {
-            writer
-                .write_all(&buffer[..len])
-                .await
-                .context("failed to write display update")?;
+            let mut fragmenter = fragmenter.context("error while encoding")?;
+            if fragmenter.size_hint() > buffer.len() {
+                buffer.resize(fragmenter.size_hint(), 0);
+            }
+
+            while let Some(len) = fragmenter.next(buffer) {
+                writer
+                    .write_all(&buffer[..len])
+                    .await
+                    .context("failed to write display update")?;
+            }
         }
 
         Ok((RunState::Continue, encoder))