Optimize accessing the foreign registry by referencing either it or the global reference in the TLS variable

Author: adamreichold

rayon-core/src/registry.rs

@@ -217,9 +217,28 @@ fn default_global_registry() -> Result<Arc<Registry>, ThreadPoolBuildError> {
     result
 }
 
-// A registry which is used only temporarily and is not associated to a worker thread.
+// A pointer to the global registry if it was initialized or a temporary reference if the current thread is not a worker thread.
 thread_local! {
-    static FOREIGN_REGISTRY: Cell<*const Arc<Registry>> = const { Cell::new(ptr::null()) };
+    static CURRENT_REGISTRY: Cell<*const Arc<Registry>> = const { Cell::new(ptr::null()) };
+}
+
+#[cold]
+fn set_current_registry_to_global_registry() -> *const Arc<Registry> {
+    let global = global_registry();
+
+    CURRENT_REGISTRY.with(|current_registry| current_registry.set(global));
+
+    global
+}
+
+unsafe fn current_registry<'a>() -> &'a Arc<Registry> {
+    let mut current = CURRENT_REGISTRY.with(Cell::get);
+
+    if current.is_null() {
+        current = set_current_registry_to_global_registry();
+    }
+
+    &*current
 }
 
 struct Terminator<'a>(&'a Arc<Registry>);
@@ -320,13 +339,7 @@ impl Registry {
         unsafe {
             let worker_thread = WorkerThread::current();
             let registry = if worker_thread.is_null() {
-                let foreign_registry = FOREIGN_REGISTRY.with(Cell::get);
-
-                if foreign_registry.is_null() {
-                    global_registry()
-                } else {
-                    &*foreign_registry
-                }
+                current_registry()
             } else {
                 &(*worker_thread).registry
             };
@@ -338,19 +351,22 @@ impl Registry {
     where
         F: FnOnce() -> R,
     {
-        struct Guard;
+        struct Guard {
+            current: *const Arc<Registry>,
+        }
 
         impl Guard {
             fn new(registry: &Arc<Registry>) -> Self {
-                FOREIGN_REGISTRY.with(|foreign_registry| foreign_registry.set(registry));
+                let current =
+                    CURRENT_REGISTRY.with(|current_registry| current_registry.replace(registry));
 
-                Self
+                Self { current }
             }
         }
 
         impl Drop for Guard {
             fn drop(&mut self) {
-                FOREIGN_REGISTRY.with(|foreign_registry| foreign_registry.set(ptr::null()));
+                CURRENT_REGISTRY.with(|current_registry| current_registry.set(self.current));
             }
         }