Bug 1577439 - Shutdown Servo's thread-pool in leak-checking builds, leak the atom table elsewhere. r=bholley

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

--HG--
extra : moz-landing-system : lando

servo/components/style/global_style_data.rs

@@ -12,6 +12,7 @@ use crate::shared_lock::SharedRwLock;
 use crate::thread_state;
 use rayon;
 use std::env;
+use parking_lot::{RwLock, RwLockReadGuard};
 
 /// Global style data
 pub struct GlobalStyleData {
@@ -22,13 +23,16 @@ pub struct GlobalStyleData {
     pub options: StyleSystemOptions,
 }
 
-/// Global thread pool
+/// Global thread pool.
 pub struct StyleThreadPool {
     /// How many threads parallel styling can use.
     pub num_threads: usize,
 
     /// The parallel styling thread pool.
-    pub style_thread_pool: Option<rayon::ThreadPool>,
+    ///
+    /// For leak-checking purposes, we want to terminate the thread-pool, which
+    /// waits for all the async jobs to complete. Thus the RwLock.
+    style_thread_pool: RwLock<Option<rayon::ThreadPool>>,
 }
 
 fn thread_name(index: usize) -> String {
@@ -57,6 +61,21 @@ fn thread_shutdown(_: usize) {
     }
 }
 
+impl StyleThreadPool {
+    /// Shuts down the thread pool, waiting for all work to complete.
+    pub fn shutdown(&self) {
+        let _ = self.style_thread_pool.write().take();
+    }
+
+    /// Returns a reference to the thread pool.
+    ///
+    /// We only really want to give read-only access to the pool, except
+    /// for shutdown().
+    pub fn pool(&self) -> RwLockReadGuard<Option<rayon::ThreadPool>> {
+        self.style_thread_pool.read()
+    }
+}
+
 lazy_static! {
     /// Global thread pool
     pub static ref STYLE_THREAD_POOL: StyleThreadPool = {
@@ -113,10 +132,11 @@ lazy_static! {
         };
 
         StyleThreadPool {
-            num_threads: num_threads,
-            style_thread_pool: pool,
+            num_threads,
+            style_thread_pool: RwLock::new(pool),
         }
     };
+
     /// Global style data
     pub static ref GLOBAL_STYLE_DATA: GlobalStyleData = GlobalStyleData {
         shared_lock: SharedRwLock::new_leaked(),

servo/ports/geckolib/glue.rs

@@ -253,8 +253,10 @@ fn traverse_subtree(
     debug!("Traversing subtree from {:?}", element);
 
     let thread_pool_holder = &*STYLE_THREAD_POOL;
+    let pool;
     let thread_pool = if traversal_flags.contains(TraversalFlags::ParallelTraversal) {
-        thread_pool_holder.style_thread_pool.as_ref()
+        pool = thread_pool_holder.pool();
+        pool.as_ref()
     } else {
         None
     };
@@ -1415,7 +1417,7 @@ pub unsafe extern "C" fn Servo_StyleSheet_FromUTF8BytesAsync(
         should_record_use_counters,
     );
 
-    if let Some(thread_pool) = STYLE_THREAD_POOL.style_thread_pool.as_ref() {
+    if let Some(thread_pool) = STYLE_THREAD_POOL.pool().as_ref() {
         thread_pool.spawn(|| {
             profiler_label!(Parse);
             async_parser.parse();
@@ -1425,6 +1427,12 @@ pub unsafe extern "C" fn Servo_StyleSheet_FromUTF8BytesAsync(
     }
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn Servo_ShutdownThreadPool() {
+    debug_assert!(is_main_thread() && !is_in_servo_traversal());
+    STYLE_THREAD_POOL.shutdown();
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn Servo_StyleSheet_FromSharedData(
     extra_data: *mut URLExtraData,

xpcom/build/XPCOMInit.cpp

@@ -88,6 +88,7 @@
 #include "mozilla/ClearOnShutdown.h"
 #include "mozilla/CountingAllocatorBase.h"
 #include "mozilla/UniquePtr.h"
+#include "mozilla/ServoStyleConsts.h"
 
 #include "mozilla/ipc/GeckoChildProcessHost.h"
 
@@ -769,7 +770,19 @@ nsresult ShutdownXPCOM(nsIServiceManager* aServMgr) {
 
   GkRust_Shutdown();
 
+#ifdef NS_FREE_PERMANENT_DATA
+  // By the time we're shutting down, there may still be async parse tasks going
+  // on in the Servo thread-pool. This is fairly uncommon, though not
+  // impossible. CSS parsing heavily uses the atom table, so obviously it's not
+  // fine to get rid of it.
+  //
+  // In leak-checking / ASAN / etc. builds, shut down the servo thread-pool,
+  // which will wait for all the work to be done. For other builds, we don't
+  // really want to wait on shutdown for possibly slow tasks. So just leak the
+  // atom table in those.
+  Servo_ShutdownThreadPool();
   NS_ShutdownAtomTable();
+#endif
 
   NS_IF_RELEASE(gDebug);