Bug 1836948 - Use the main thread as part of the style thread pool. r=smaug

Differential Revision: https://phabricator.services.mozilla.com/D180079

Cargo.lock

@@ -5196,6 +5196,7 @@ dependencies = [
  "parking_lot",
  "precomputed-hash",
  "rayon",
+ "rayon-core",
  "regex",
  "selectors",
  "serde",

servo/components/style/Cargo.toml

@@ -57,6 +57,7 @@ owning_ref = "0.4"
 parking_lot = "0.12"
 precomputed-hash = "0.1.1"
 rayon = "1"
+rayon-core = "1"
 selectors = { path = "../selectors" }
 serde = {version = "1.0", optional = true, features = ["derive"]}
 servo_arc = { path = "../servo_arc" }

servo/components/style/driver.rs

@@ -15,7 +15,6 @@ use crate::scoped_tls::ScopedTLS;
 use crate::traversal::{DomTraversal, PerLevelTraversalData, PreTraverseToken};
 use rayon;
 use std::collections::VecDeque;
-use std::mem;
 use time;
 
 #[cfg(feature = "servo")]
@@ -106,17 +105,19 @@ where
     //     ThreadLocalStyleContext on the main thread. If the main thread
     //     ThreadLocalStyleContext has not released its TLS borrow by that point,
     //     we'll panic on double-borrow.
-    let mut scoped_tls = pool.map(ScopedTLS::<ThreadLocalStyleContext<E>>::new);
-    let mut tlc = ThreadLocalStyleContext::new();
-    let mut context = StyleContext {
-        shared: traversal.shared_context(),
-        thread_local: &mut tlc,
-    };
-
+    let mut scoped_tls = ScopedTLS::<ThreadLocalStyleContext<E>>::new(pool);
     // Process the nodes breadth-first. This helps keep similar traversal characteristics for the
     // style sharing cache.
     let work_unit_max = work_unit_max();
     with_pool_in_place_scope(work_unit_max, pool, |maybe_scope| {
+        let mut tlc = scoped_tls.ensure(parallel::create_thread_local_context);
+        let mut context = StyleContext {
+            shared: traversal.shared_context(),
+            thread_local: &mut tlc,
+        };
+
+        debug_assert_eq!(scoped_tls.current_thread_index(), 0, "Main thread should be the first thread");
+
         let mut discovered = VecDeque::with_capacity(work_unit_max * 2);
         discovered.push_back(unsafe { SendNode::new(root.as_node()) });
         parallel::style_trees(
@@ -124,23 +125,19 @@ where
             discovered,
             root.as_node().opaque(),
             work_unit_max,
-            static_prefs::pref!("layout.css.stylo-local-work-queue.in-main-thread") as usize,
             PerLevelTraversalData { current_dom_depth: root.depth() },
             maybe_scope,
             traversal,
-            scoped_tls.as_ref(),
+            &scoped_tls,
         );
     });
 
     // Collect statistics from thread-locals if requested.
     if dump_stats || report_stats {
-        let mut aggregate = mem::replace(&mut context.thread_local.statistics, Default::default());
-        let parallel = pool.is_some();
-        if let Some(ref mut tls) = scoped_tls {
-            for slot in tls.slots() {
-                if let Some(cx) = slot.get_mut() {
-                    aggregate += cx.statistics.clone();
-                }
+        let mut aggregate = PerThreadTraversalStatistics::default();
+        for slot in scoped_tls.slots() {
+            if let Some(cx) = slot.get_mut() {
+                aggregate += cx.statistics.clone();
             }
         }
 
@@ -149,6 +146,7 @@ where
         }
         // dump statistics to stdout if requested
         if dump_stats {
+            let parallel = pool.is_some();
             let stats =
                 TraversalStatistics::new(aggregate, traversal, parallel, start_time.unwrap());
             if stats.is_large {

servo/components/style/global_style_data.rs

@@ -112,6 +112,9 @@ impl StyleThreadPool {
         while let Some(join_handle) = STYLE_THREAD_JOIN_HANDLES.lock().pop() {
             let _ = join_handle.join();
         }
+
+        // Clean up the current thread too.
+        rayon_core::clean_up_use_current_thread();
     }
 
     /// Returns a reference to the thread pool.
@@ -172,11 +175,14 @@ lazy_static! {
         };
 
         let num_threads = cmp::min(num_threads, STYLO_MAX_THREADS);
-        let (pool, num_threads) = if num_threads < 1 {
+        // Since the main-thread is also part of the pool, having one thread or less doesn't make
+        // sense.
+        let (pool, num_threads) = if num_threads <= 1 {
             (None, None)
         } else {
             let workers = rayon::ThreadPoolBuilder::new()
                 .spawn_handler(thread_spawn)
+                .use_current_thread()
                 .num_threads(num_threads)
                 .thread_name(thread_name)
                 .start_handler(thread_startup)

servo/components/style/parallel.rs

@@ -54,7 +54,7 @@ pub const STACK_SAFETY_MARGIN_KB: usize = 168;
 /// out of line so we don't allocate stack space for the entire struct
 /// in the caller.
 #[inline(never)]
-fn create_thread_local_context<'scope, E>(slot: &mut Option<ThreadLocalStyleContext<E>>)
+pub (crate) fn create_thread_local_context<'scope, E>(slot: &mut Option<ThreadLocalStyleContext<E>>)
 where
     E: TElement + 'scope,
 {
@@ -86,11 +86,10 @@ fn distribute_one_chunk<'a, 'scope, E, D>(
             items,
             traversal_root,
             work_unit_max,
-            static_prefs::pref!("layout.css.stylo-local-work-queue.in-worker") as usize,
             traversal_data,
             Some(scope),
             traversal,
-            Some(tls),
+            tls,
         );
     })
 }
@@ -139,15 +138,20 @@ pub fn style_trees<'a, 'scope, E, D>(
     mut discovered: VecDeque<SendNode<E::ConcreteNode>>,
     traversal_root: OpaqueNode,
     work_unit_max: usize,
-    local_queue_size: usize,
     mut traversal_data: PerLevelTraversalData,
     scope: Option<&'a rayon::ScopeFifo<'scope>>,
     traversal: &'scope D,
-    tls: Option<&'scope ScopedTLS<'scope, ThreadLocalStyleContext<E>>>,
+    tls: &'scope ScopedTLS<'scope, ThreadLocalStyleContext<E>>,
 ) where
     E: TElement + 'scope,
     D: DomTraversal<E>,
 {
+    let local_queue_size = if tls.current_thread_index() == 0 {
+        static_prefs::pref!("layout.css.stylo-local-work-queue.in-main-thread")
+    } else {
+        static_prefs::pref!("layout.css.stylo-local-work-queue.in-worker")
+    } as usize;
+
     let mut nodes_remaining_at_current_depth = discovered.len();
     while let Some(node) = discovered.pop_front() {
         let mut children_to_process = 0isize;
@@ -178,7 +182,7 @@ pub fn style_trees<'a, 'scope, E, D>(
                 traversal_data_copy,
                 scope.unwrap(),
                 traversal,
-                tls.unwrap(),
+                tls,
             );
         }
 

servo/components/style/scoped_tls.rs

@@ -20,7 +20,7 @@ use std::ops::DerefMut;
 /// Note that the cleanup is done on the thread that owns the scoped TLS, thus
 /// the Send bound.
 pub struct ScopedTLS<'scope, T: Send> {
-    pool: &'scope rayon::ThreadPool,
+    pool: Option<&'scope rayon::ThreadPool>,
     slots: [RefCell<Option<T>>; STYLO_MAX_THREADS],
 }
 
@@ -31,23 +31,29 @@ unsafe impl<'scope, T: Send> Sync for ScopedTLS<'scope, T> {}
 impl<'scope, T: Send> ScopedTLS<'scope, T> {
     /// Create a new scoped TLS that will last as long as this rayon threadpool
     /// reference.
-    pub fn new(pool: &'scope rayon::ThreadPool) -> Self {
-        debug_assert!(pool.current_num_threads() <= STYLO_MAX_THREADS);
+    pub fn new(pool: Option<&'scope rayon::ThreadPool>) -> Self {
+        debug_assert!(pool.map_or(true, |p| p.current_num_threads() <= STYLO_MAX_THREADS));
         ScopedTLS {
             pool,
             slots: Default::default(),
         }
     }
 
+    /// Returns the index corresponding to the calling thread in the thread pool.
+    #[inline]
+    pub fn current_thread_index(&self) -> usize {
+        self.pool.map_or(0, |p| p.current_thread_index().unwrap())
+    }
+
     /// Return an immutable reference to the `Option<T>` that this thread owns.
     pub fn borrow(&self) -> Ref<Option<T>> {
-        let idx = self.pool.current_thread_index().unwrap();
+        let idx = self.current_thread_index();
         self.slots[idx].borrow()
     }
 
     /// Return a mutable reference to the `Option<T>` that this thread owns.
     pub fn borrow_mut(&self) -> RefMut<Option<T>> {
-        let idx = self.pool.current_thread_index().unwrap();
+        let idx = self.current_thread_index();
         self.slots[idx].borrow_mut()
     }