moss/client: Blit entries in parallel

As children of dirs are pre-structured by vfs, we can blit them all in parallel safely.
AerynOS · Feb 25, 2025 · 12f0388 · 12f0388
1 parent 5e85a62
commit 12f0388
Showing 1 changed file with 33 additions and 16 deletions.
diff --git a/moss/src/client/mod.rs b/moss/src/client/mod.rs
@@ -13,6 +13,7 @@ use std::{
     fmt, io,
     os::{fd::RawFd, unix::fs::symlink},
     path::{Path, PathBuf},
+    sync::{Arc, RwLock},
     time::{Duration, Instant},
 };
 
@@ -26,6 +27,7 @@ use nix::{
     unistd::{close, linkat, mkdir, symlinkat},
 };
 use postblit::TriggerScope;
+use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use stone::{payload::layout, read::PayloadKind};
 use thiserror::Error;
 use tui::{MultiProgress, ProgressBar, ProgressStyle, Styled};
@@ -625,7 +627,7 @@ impl Client {
         progress.tick();
 
         let now = Instant::now();
-        let mut stats = BlitStats::default();
+        let stats = Arc::new(RwLock::new(BlitStats::default()));
 
         let tree = self.vfs(packages)?;
 
@@ -648,9 +650,9 @@ impl Client {
             let root_dir = fcntl::open(&blit_target, OFlag::O_DIRECTORY | OFlag::O_RDONLY, Mode::empty())?;
 
             if let Element::Directory(_, _, children) = root {
-                for child in children {
-                    self.blit_element(root_dir, cache_fd, child, &progress, &mut stats)?;
-                }
+                children
+                    .par_iter()
+                    .try_for_each(|child| self.blit_element(root_dir, cache_fd, child, &progress, &stats))?;
             }
 
             close(root_dir)?;
@@ -659,7 +661,8 @@ impl Client {
         progress.finish_and_clear();
 
         let elapsed = now.elapsed();
-        let num_entries = stats.num_entries();
+        let stats_raw = stats.write().unwrap();
+        let num_entries = stats_raw.num_entries();
 
         println!(
             "\n{} entries blitted in {} {}",
@@ -678,9 +681,9 @@ impl Client {
         &self,
         parent: RawFd,
         cache: RawFd,
-        element: Element<'_, PendingFile>,
+        element: &Element<'_, PendingFile>,
         progress: &ProgressBar,
-        stats: &mut BlitStats,
+        stats: &Arc<RwLock<BlitStats>>,
     ) -> Result<(), Error> {
         progress.inc(1);
         match element {
@@ -689,10 +692,15 @@ impl Client {
                 self.blit_element_item(parent, cache, name, item, stats)?;
 
                 // open the new dir
-                let newdir = fcntl::openat(parent, name, OFlag::O_RDONLY | OFlag::O_DIRECTORY, Mode::empty())?;
-                for child in children.into_iter() {
-                    self.blit_element(newdir, cache, child, progress, stats)?;
-                }
+                let newdir = fcntl::openat(
+                    parent,
+                    name.to_owned(),
+                    OFlag::O_RDONLY | OFlag::O_DIRECTORY,
+                    Mode::empty(),
+                )?;
+                children
+                    .par_iter()
+                    .try_for_each(|child| self.blit_element(newdir, cache, child, progress, stats))?;
                 close(newdir)?;
                 Ok(())
             }
@@ -717,7 +725,7 @@ impl Client {
         cache: RawFd,
         subpath: &str,
         item: &PendingFile,
-        stats: &mut BlitStats,
+        stats: &Arc<RwLock<BlitStats>>,
     ) -> Result<(), Error> {
         match &item.layout.entry {
             layout::Entry::Regular(id, _) => {
@@ -763,15 +771,24 @@ impl Client {
                     }
                 }
 
-                stats.num_files += 1;
+                let mut stats_raw = stats.write().unwrap();
+                {
+                    stats_raw.num_files += 1;
+                }
             }
             layout::Entry::Symlink(source, _) => {
                 symlinkat(source.as_str(), Some(parent), subpath)?;
-                stats.num_symlinks += 1;
+                let mut stats_raw = stats.write().unwrap();
+                {
+                    stats_raw.num_symlinks += 1;
+                }
             }
             layout::Entry::Directory(_) => {
                 mkdirat(parent, subpath, Mode::from_bits_truncate(item.layout.mode))?;
-                stats.num_dirs += 1;
+                let mut stats_raw = stats.write().unwrap();
+                {
+                    stats_raw.num_dirs += 1;
+                }
             }
 
             // unimplemented
@@ -975,7 +992,7 @@ struct BlitStats {
     num_dirs: u64,
 }
 
-impl BlitStats {
+impl<'a> BlitStats {
     fn num_entries(&self) -> u64 {
         self.num_files + self.num_symlinks + self.num_dirs
     }