Port with_thread_data from word_lock into parking_lot

Author: faern

core/src/parking_lot.rs

@@ -162,8 +162,11 @@ impl ThreadData {
     }
 }
 
-// Returns a ThreadData structure for the current thread
-unsafe fn get_thread_data(local: &mut Option<ThreadData>) -> &ThreadData {
+// Invokes the given closure with a reference to the current thread `ThreadData`.
+fn with_thread_data<F, T>(f: F) -> T
+where
+    F: FnOnce(&ThreadData) -> T,
+{
     // Try to read from thread-local storage, but return None if the TLS has
     // already been destroyed.
     #[cfg(has_localkey_try_with)]
@@ -175,16 +178,22 @@ unsafe fn get_thread_data(local: &mut Option<ThreadData>) -> &ThreadData {
         panic::catch_unwind(|| key.with(|x| x as *const ThreadData)).ok()
     }
 
-    // Unlike word_lock::ThreadData, parking_lot::ThreadData is always expensive
-    // to construct. Try to use a thread-local version if possible.
-    thread_local!(static THREAD_DATA: ThreadData = ThreadData::new());
-    if let Some(tls) = try_get_tls(&THREAD_DATA) {
-        return &*tls;
+    let mut thread_data_ptr = ptr::null();
+    // If ThreadData is expensive to construct, then we want to use a cached
+    // version in thread-local storage if possible.
+    if !ThreadParker::IS_CHEAP_TO_CONSTRUCT {
+        thread_local!(static THREAD_DATA: ThreadData = ThreadData::new());
+        if let Some(tls_thread_data) = try_get_tls(&THREAD_DATA) {
+            thread_data_ptr = tls_thread_data;
+        }
     }
-
     // Otherwise just create a ThreadData on the stack
-    *local = Some(ThreadData::new());
-    local.as_ref().unwrap()
+    let mut thread_data_storage = None;
+    if thread_data_ptr.is_null() {
+        thread_data_ptr = thread_data_storage.get_or_insert_with(ThreadData::new);
+    }
+
+    f(unsafe { &*thread_data_ptr })
 }
 
 impl Drop for ThreadData {
@@ -579,106 +588,106 @@ unsafe fn park_internal(
     timeout: Option<Instant>,
 ) -> ParkResult {
     // Grab our thread data, this also ensures that the hash table exists
-    let mut thread_data = None;
-    let thread_data = get_thread_data(&mut thread_data);
+    with_thread_data(|thread_data| {
 
-    // Lock the bucket for the given key
-    let bucket = lock_bucket(key);
+        // Lock the bucket for the given key
+        let bucket = lock_bucket(key);
 
-    // If the validation function fails, just return
-    if !validate() {
+        // If the validation function fails, just return
+        if !validate() {
+            bucket.mutex.unlock();
+            return ParkResult::Invalid;
+        }
+
+        // Append our thread data to the queue and unlock the bucket
+        thread_data.parked_with_timeout.set(timeout.is_some());
+        thread_data.next_in_queue.set(ptr::null());
+        thread_data.key.store(key, Ordering::Relaxed);
+        thread_data.park_token.set(park_token);
+        thread_data.parker.prepare_park();
+        if !bucket.queue_head.get().is_null() {
+            (*bucket.queue_tail.get()).next_in_queue.set(thread_data);
+        } else {
+            bucket.queue_head.set(thread_data);
+        }
+        bucket.queue_tail.set(thread_data);
         bucket.mutex.unlock();
-        return ParkResult::Invalid;
-    }
 
-    // Append our thread data to the queue and unlock the bucket
-    thread_data.parked_with_timeout.set(timeout.is_some());
-    thread_data.next_in_queue.set(ptr::null());
-    thread_data.key.store(key, Ordering::Relaxed);
-    thread_data.park_token.set(park_token);
-    thread_data.parker.prepare_park();
-    if !bucket.queue_head.get().is_null() {
-        (*bucket.queue_tail.get()).next_in_queue.set(thread_data);
-    } else {
-        bucket.queue_head.set(thread_data);
-    }
-    bucket.queue_tail.set(thread_data);
-    bucket.mutex.unlock();
+        // Invoke the pre-sleep callback
+        before_sleep();
+
+        // Park our thread and determine whether we were woken up by an unpark or by
+        // our timeout. Note that this isn't precise: we can still be unparked since
+        // we are still in the queue.
+        let unparked = match timeout {
+            Some(timeout) => thread_data.parker.park_until(timeout),
+            None => {
+                thread_data.parker.park();
+                // call deadlock detection on_unpark hook
+                deadlock::on_unpark(thread_data);
+                true
+            }
+        };
 
-    // Invoke the pre-sleep callback
-    before_sleep();
-
-    // Park our thread and determine whether we were woken up by an unpark or by
-    // our timeout. Note that this isn't precise: we can still be unparked since
-    // we are still in the queue.
-    let unparked = match timeout {
-        Some(timeout) => thread_data.parker.park_until(timeout),
-        None => {
-            thread_data.parker.park();
-            // call deadlock detection on_unpark hook
-            deadlock::on_unpark(thread_data);
-            true
+        // If we were unparked, return now
+        if unparked {
+            return ParkResult::Unparked(thread_data.unpark_token.get());
         }
-    };
-
-    // If we were unparked, return now
-    if unparked {
-        return ParkResult::Unparked(thread_data.unpark_token.get());
-    }
 
-    // Lock our bucket again. Note that the hashtable may have been rehashed in
-    // the meantime. Our key may also have changed if we were requeued.
-    let (key, bucket) = lock_bucket_checked(&thread_data.key);
+        // Lock our bucket again. Note that the hashtable may have been rehashed in
+        // the meantime. Our key may also have changed if we were requeued.
+        let (key, bucket) = lock_bucket_checked(&thread_data.key);
 
-    // Now we need to check again if we were unparked or timed out. Unlike the
-    // last check this is precise because we hold the bucket lock.
-    if !thread_data.parker.timed_out() {
-        bucket.mutex.unlock();
-        return ParkResult::Unparked(thread_data.unpark_token.get());
-    }
+        // Now we need to check again if we were unparked or timed out. Unlike the
+        // last check this is precise because we hold the bucket lock.
+        if !thread_data.parker.timed_out() {
+            bucket.mutex.unlock();
+            return ParkResult::Unparked(thread_data.unpark_token.get());
+        }
 
-    // We timed out, so we now need to remove our thread from the queue
-    let mut link = &bucket.queue_head;
-    let mut current = bucket.queue_head.get();
-    let mut previous = ptr::null();
-    while !current.is_null() {
-        if current == thread_data {
-            let next = (*current).next_in_queue.get();
-            link.set(next);
-            let mut was_last_thread = true;
-            if bucket.queue_tail.get() == current {
-                bucket.queue_tail.set(previous);
-            } else {
-                // Scan the rest of the queue to see if there are any other
-                // entries with the given key.
-                let mut scan = next;
-                while !scan.is_null() {
-                    if (*scan).key.load(Ordering::Relaxed) == key {
-                        was_last_thread = false;
-                        break;
+        // We timed out, so we now need to remove our thread from the queue
+        let mut link = &bucket.queue_head;
+        let mut current = bucket.queue_head.get();
+        let mut previous = ptr::null();
+        while !current.is_null() {
+            if current == thread_data {
+                let next = (*current).next_in_queue.get();
+                link.set(next);
+                let mut was_last_thread = true;
+                if bucket.queue_tail.get() == current {
+                    bucket.queue_tail.set(previous);
+                } else {
+                    // Scan the rest of the queue to see if there are any other
+                    // entries with the given key.
+                    let mut scan = next;
+                    while !scan.is_null() {
+                        if (*scan).key.load(Ordering::Relaxed) == key {
+                            was_last_thread = false;
+                            break;
+                        }
+                        scan = (*scan).next_in_queue.get();
                     }
-                    scan = (*scan).next_in_queue.get();
                 }
-            }
 
-            // Callback to indicate that we timed out, and whether we were the
-            // last thread on the queue.
-            timed_out(key, was_last_thread);
-            break;
-        } else {
-            link = &(*current).next_in_queue;
-            previous = current;
-            current = link.get();
+                // Callback to indicate that we timed out, and whether we were the
+                // last thread on the queue.
+                timed_out(key, was_last_thread);
+                break;
+            } else {
+                link = &(*current).next_in_queue;
+                previous = current;
+                current = link.get();
+            }
         }
-    }
 
-    // There should be no way for our thread to have been removed from the queue
-    // if we timed out.
-    debug_assert!(!current.is_null());
+        // There should be no way for our thread to have been removed from the queue
+        // if we timed out.
+        debug_assert!(!current.is_null());
 
-    // Unlock the bucket, we are done
-    bucket.mutex.unlock();
-    ParkResult::TimedOut
+        // Unlock the bucket, we are done
+        bucket.mutex.unlock();
+        ParkResult::TimedOut
+    })
 }
 
 /// Unparks one thread from the queue associated with the given key.
@@ -1149,7 +1158,7 @@ pub mod deadlock {
 
 #[cfg(feature = "deadlock_detection")]
 mod deadlock_impl {
-    use super::{get_hashtable, get_thread_data, lock_bucket, ThreadData, NUM_THREADS};
+    use super::{get_hashtable, with_thread_data, lock_bucket, ThreadData, NUM_THREADS};
     use backtrace::Backtrace;
     use petgraph;
     use petgraph::graphmap::DiGraphMap;
@@ -1222,19 +1231,19 @@ mod deadlock_impl {
     }
 
     pub unsafe fn acquire_resource(key: usize) {
-        let mut thread_data = None;
-        let thread_data = get_thread_data(&mut thread_data);
-        (*thread_data.deadlock_data.resources.get()).push(key);
+        with_thread_data(|thread_data| {
+            (*thread_data.deadlock_data.resources.get()).push(key);
+        });
     }
 
     pub unsafe fn release_resource(key: usize) {
-        let mut thread_data = None;
-        let thread_data = get_thread_data(&mut thread_data);
-        let resources = &mut (*thread_data.deadlock_data.resources.get());
-        match resources.iter().rposition(|x| *x == key) {
-            Some(p) => resources.swap_remove(p),
-            None => panic!("key {} not found in thread resources", key),
-        };
+        with_thread_data(|thread_data| {
+            let resources = &mut (*thread_data.deadlock_data.resources.get());
+            match resources.iter().rposition(|x| *x == key) {
+                Some(p) => resources.swap_remove(p),
+                None => panic!("key {} not found in thread resources", key),
+            };
+        });
     }
 
     pub fn check_deadlock() -> Vec<Vec<DeadlockedThread>> {