Share persistent worker bins and caches

src/context.rs

@@ -43,6 +43,7 @@ use crate::linker_script::{DynamicPattern, VersionPattern};
 use crate::mapped_file::MappedFile;
 use crate::symbol::{Bins, Symbol, SymbolChunkId, SymbolId, SymbolSlot, SymbolTable};
 use crate::util::perf::Timers;
+use crate::util::worker_local::WorkerLocal;
 
 // Keep immutable and mutable chunk lookup in the same static match.
 macro_rules! chunk_header {
@@ -160,7 +161,7 @@ pub struct Context<E: Arch> {
 
     /// Global symbol keys recorded while input files are parsed, one bin per
     /// Rayon worker plus one for callers outside the pool.
-    symbol_bins: OnceLock<Vec<Mutex<Bins<SymbolSlot>>>>,
+    symbol_bins: OnceLock<WorkerLocal<Bins<SymbolSlot>>>,
 
     pub objs: FileList<ObjectFile<E>>,
     pub dsos: FileList<SharedFile<E>>,
@@ -363,15 +364,9 @@ impl<E: Arch> Context<E> {
     /// Returns this worker's symbol bin. Looking it up once per file keeps the
     /// synchronization cost outside the per-symbol loop.
     pub(crate) fn symbol_bin(&self) -> MutexGuard<'_, Bins<SymbolSlot>> {
-        let bins = self.symbol_bins.get_or_init(|| {
-            let workers = rayon::current_num_threads();
-            (0..=workers).map(|_| Mutex::new(Bins::new())).collect()
-        });
-        let fallback = bins.len() - 1;
-        let index = rayon::current_thread_index()
-            .unwrap_or(fallback)
-            .min(fallback);
-        bins[index].lock().unwrap()
+        self.symbol_bins
+            .get_or_init(|| WorkerLocal::new(Bins::new))
+            .get()
     }
 
     /// Takes all keys recorded since the previous gather, leaving empty bins
@@ -379,11 +374,7 @@ impl<E: Arch> Context<E> {
     pub(crate) fn take_symbol_bins(&mut self) -> Vec<Bins<SymbolSlot>> {
         self.symbol_bins
             .get_mut()
-            .map(|bins| {
-                bins.iter_mut()
-                    .map(|bin| std::mem::take(bin.get_mut().unwrap()))
-                    .collect()
-            })
+            .map(|bins| bins.take().collect())
             .unwrap_or_default()
     }
 

src/passes.rs

@@ -1,5 +1,6 @@
 //! The passes of a link, in roughly the order the driver runs them.
 
+use crate::util::worker_local::WorkerLocal;
 use std::borrow::Cow;
 use std::cell::UnsafeCell;
 use std::collections::{HashMap, HashSet};
@@ -548,13 +549,9 @@ fn parse_input_sections<E: Arch>(ctx: &mut Context<E>) {
     let (bins, pending): (Vec<Bins<ComdatSymbolSlot>>, Vec<Vec<PendingComdatOwner>>) = {
         let Context { objs, symbols, .. } = ctx;
         // Reuse each worker's buffers across jobs, locking once per file.
-        let workers = rayon::current_num_threads();
-        let work: Vec<_> = (0..=workers)
-            .map(|_| Mutex::new((Bins::new(), Vec::new())))
-            .collect();
+        let work = WorkerLocal::new(|| (Bins::new(), Vec::new()));
         objs.par_iter_mut().for_each(|file| {
-            let worker = rayon::current_thread_index().unwrap_or(workers);
-            let mut local = work[worker].lock().unwrap();
+            let mut local = work.get();
             let (bins, pending) = &mut *local;
             if file.base.is_reachable() {
                 if file.base.mf.is_some() && !file.is_lto_input() && !file.sections_parsed {
@@ -590,9 +587,7 @@ fn parse_input_sections<E: Arch>(ctx: &mut Context<E>) {
                 }
             }
         });
-        work.into_iter()
-            .map(|bin| bin.into_inner().unwrap())
-            .unzip()
+        work.into_values().unzip()
     };
 
     drop(t);
@@ -951,14 +946,9 @@ pub fn create_merged_sections<E: Arch>(ctx: &mut Context<E>) {
         } = ctx;
         let merged = RwLock::new(std::mem::take(merged_sections));
         // Keep a cache for each actual worker, including across Rayon jobs.
-        // Padding keeps independent workers' mutexes off the same cache line.
-        #[repr(align(128))]
-        struct Cache(Mutex<crate::chunks::merged::MergedSectionCache>);
-        let workers = rayon::current_num_threads();
-        let caches: Vec<_> = (0..=workers).map(|_| Cache(Mutex::default())).collect();
+        let caches = WorkerLocal::new(crate::chunks::merged::MergedSectionCache::default);
         objs.par_iter_mut().for_each(|file| {
-            let worker = rayon::current_thread_index().unwrap_or(workers);
-            let mut cache = caches[worker].0.lock().unwrap();
+            let mut cache = caches.get();
             file.convert_mergeable_sections(args, &merged, &mut cache);
         });
         *merged_sections = merged.into_inner().unwrap();
@@ -1280,15 +1270,13 @@ pub fn create_output_sections<E: Arch>(ctx: &mut Context<E>) {
     let num_files = ctx.objs.len();
     let shared: Mutex<OutputSectionShared<E>> =
         Mutex::new((HashMap::new(), std::mem::take(&mut ctx.output_sections)));
-    let workers = rayon::current_num_threads();
-    let caches: Vec<_> = (0..=workers).map(|_| Mutex::new(HashMap::new())).collect();
+    let caches = WorkerLocal::new(HashMap::new);
 
     // Instantiate output sections and assign input sections to them
     {
         let Context { objs, args, .. } = ctx;
         objs.par_iter_mut().enumerate().for_each(|(fi, file)| {
-            let worker = rayon::current_thread_index().unwrap_or(workers);
-            let mut cache = caches[worker].lock().unwrap();
+            let mut cache = caches.get();
             let file_id = file.id();
             let shstrtab = file.base.shstrtab;
             let num_elf_sections = file.num_elf_sections;

src/reader.rs

@@ -9,7 +9,6 @@
 //! the results are sorted back into command line order.
 
 use std::path::Path;
-use std::sync::Mutex;
 
 use rayon::prelude::*;
 
@@ -22,6 +21,7 @@ use crate::input_files::{DsoId, FileId, ObjId, ObjectFile, SharedFile};
 use crate::linker_script::Script;
 use crate::mapped_file::{must_open_file, open_file, MappedFile};
 use crate::util::perf::Counter;
+use crate::util::worker_local::WorkerLocal;
 use crate::{fatal, out, warn};
 
 /// A file that has been read, with its command line position.
@@ -30,12 +30,6 @@ enum Loaded<E: Arch> {
     Dso(Vec<u32>, Box<SharedFile<E>>),
 }
 
-/// Appends a result to the current worker's bin.
-fn push_to_worker<T>(bins: &[Mutex<Vec<T>>], value: T) {
-    let worker = rayon::current_thread_index().unwrap_or(bins.len() - 1);
-    bins[worker].lock().unwrap().push(value);
-}
-
 fn get_file_type<E: Arch>(ctx: &Context<E>, mf: &MappedFile) -> FileType {
     filetype::get_file_type(&ctx.args.plugin, mf)
 }
@@ -91,7 +85,10 @@ fn new_shared_file<E: Arch>(
     mf: &'static MappedFile,
 ) -> SharedFile<E> {
     if rctx.is_static {
-        fatal!("{}: attempted static link of a dynamic object", mf.name.display());
+        fatal!(
+            "{}: attempted static link of a dynamic object",
+            mf.name.display()
+        );
     }
     check_machine_type(ctx, mf);
     let mut file = SharedFile::<E>::new(mf, &mut ctx.symbol_bin());
@@ -251,7 +248,8 @@ fn open_library<E: Arch>(
         if target != E::NAME {
             warn!(
                 "{}: skipping incompatible file: {target} (e_machine {})",
-                path.display(), E::E_MACHINE
+                path.display(),
+                E::E_MACHINE
             );
             return None;
         }
@@ -337,9 +335,8 @@ pub fn read_input_files<E: Arch>(ctx: &mut Context<E>, jobs: Vec<ReaderJob>) {
 
     // Keep one bin per worker, locking only to append a result after I/O and
     // parsing. Command line positions restore the order after the jobs join.
-    let workers = rayon::current_num_threads();
-    let scripts: Vec<_> = (0..=workers).map(|_| Mutex::new(Vec::new())).collect();
-    let loaded: Vec<_> = (0..=workers).map(|_| Mutex::new(Vec::new())).collect();
+    let scripts = WorkerLocal::new(Vec::new);
+    let loaded = WorkerLocal::new(Vec::new);
 
     let ctx_ref: &Context<E> = ctx;
     rayon::scope(|scope| {
@@ -373,7 +370,7 @@ pub fn read_input_files<E: Arch>(ctx: &mut Context<E>, jobs: Vec<ReaderJob>) {
                             if let Some(file) =
                                 read_archive_member(ctx_ref, child_rctx, child, archive_name)
                             {
-                                push_to_worker(loaded, file);
+                                loaded.get().push(file);
                             }
                         });
                     }
@@ -388,19 +385,19 @@ pub fn read_input_files<E: Arch>(ctx: &mut Context<E>, jobs: Vec<ReaderJob>) {
                             if let Some(file) =
                                 read_archive_member(ctx_ref, child_rctx, child, archive_name)
                             {
-                                push_to_worker(loaded, file);
+                                loaded.get().push(file);
                             }
                         });
                     }
                 }
-                FileType::Text => push_to_worker(&scripts, (rctx, mf)),
+                FileType::Text => scripts.get().push((rctx, mf)),
                 FileType::ElfObj => {
                     let file = new_object_file(ctx_ref, &rctx, mf, Path::new(""));
-                    push_to_worker(&loaded, Loaded::Obj(rctx.pos, Box::new(file)));
+                    loaded.get().push(Loaded::Obj(rctx.pos, Box::new(file)));
                 }
                 FileType::ElfDso => {
                     let file = new_shared_file(ctx_ref, &rctx, mf);
-                    push_to_worker(&loaded, Loaded::Dso(rctx.pos, Box::new(file)));
+                    loaded.get().push(Loaded::Dso(rctx.pos, Box::new(file)));
                 }
                 FileType::GccLtoObj | FileType::LlvmBitcode => {
                     defer_lto_object(ctx_ref, rctx, mf, Path::new(""));
@@ -410,15 +407,12 @@ pub fn read_input_files<E: Arch>(ctx: &mut Context<E>, jobs: Vec<ReaderJob>) {
         });
     });
 
-    for l in loaded.into_iter().flat_map(|bin| bin.into_inner().unwrap()) {
+    for l in loaded.into_values().flatten() {
         push_loaded(ctx, l);
     }
 
     // Parse linker scripts and read the files they name.
-    let mut scripts: Vec<_> = scripts
-        .into_iter()
-        .flat_map(|bin| bin.into_inner().unwrap())
-        .collect();
+    let mut scripts: Vec<_> = scripts.into_values().flatten().collect();
     scripts.sort_by(|(a, _), (b, _)| a.pos.cmp(&b.pos));
     for (mut rctx, mf) in scripts {
         Script::new(ctx, &mut rctx, mf).parse_linker_script();

src/util/mod.rs

@@ -13,6 +13,7 @@ pub mod perf;
 mod prefetch;
 pub(crate) mod siphash;
 pub mod tar;
+pub(crate) mod worker_local;
 
 pub(crate) use prefetch::prefetch;
 

src/util/worker_local.rs

@@ -0,0 +1,70 @@
+use std::sync::{Mutex, MutexGuard};
+
+// Keep independent workers' locks and buffers off the same cache line.
+#[repr(align(128))]
+struct Worker<T>(Mutex<T>);
+
+/// Persistent scratch storage per Rayon worker, plus a slot for callers
+/// outside the pool. Look up and lock once per file, not once per symbol.
+pub(crate) struct WorkerLocal<T>(Vec<Worker<T>>);
+
+impl<T> WorkerLocal<T> {
+    pub fn new(init: impl Fn() -> T) -> Self {
+        Self(
+            (0..=rayon::current_num_threads())
+                .map(|_| Worker(Mutex::new(init())))
+                .collect(),
+        )
+    }
+
+    pub fn get(&self) -> MutexGuard<'_, T> {
+        let fallback = self.0.len() - 1;
+        let worker = rayon::current_thread_index()
+            .unwrap_or(fallback)
+            .min(fallback);
+        self.0[worker].0.lock().unwrap()
+    }
+
+    pub fn into_values(self) -> impl Iterator<Item = T> {
+        self.0
+            .into_iter()
+            .map(|worker| worker.0.into_inner().unwrap())
+    }
+
+    pub fn take(&mut self) -> impl Iterator<Item = T> + '_
+    where
+        T: Default,
+    {
+        self.0
+            .iter_mut()
+            .map(|worker| std::mem::take(worker.0.get_mut().unwrap()))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rayon::prelude::*;
+
+    #[test]
+    fn reuses_worker_storage_and_keeps_an_outside_slot() {
+        let pool = rayon::ThreadPoolBuilder::new()
+            .num_threads(2)
+            .build()
+            .unwrap();
+        let mut slots = pool.install(|| {
+            let slots = WorkerLocal::new(Vec::new);
+            for batch in 0..3 {
+                (0..100)
+                    .into_par_iter()
+                    .for_each(|i| slots.get().push(batch * 100 + i));
+            }
+            slots
+        });
+        slots.get().push(300);
+        let mut values: Vec<_> = slots.take().flatten().collect();
+        values.sort_unstable();
+        assert_eq!(values, (0..301).collect::<Vec<_>>());
+        assert!(slots.into_values().all(|slot| slot.is_empty()));
+    }
+}