Handle synchronization in a private module with safer closure

components/profile/heartbeats.rs

@@ -2,61 +2,37 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
-
 use heartbeats_simple::HeartbeatPow as Heartbeat;
-use heartbeats_simple::HeartbeatPowContext as HeartbeatContext;
 use profile_traits::time::ProfilerCategory;
+use self::synchronized_heartbeat::{heartbeat_window_callback, lock_and_work};
 use servo_config::opts;
 use std::collections::HashMap;
 use std::env::var_os;
 use std::error::Error;
 use std::fs::File;
 use std::path::Path;
-use std::sync::atomic::{ATOMIC_BOOL_INIT, AtomicBool, Ordering};
-
-
-static mut HBS: Option<*mut HashMap<ProfilerCategory, Heartbeat>> = None;
-
-// unfortunately can't encompass the actual hashmap in a Mutex (Heartbeat isn't Send/Sync), so we'll use a spinlock
-static mut HBS_SPINLOCK: AtomicBool = ATOMIC_BOOL_INIT;
-
-fn lock_and_work<F, R>(work: F) -> R where F: FnOnce() -> R {
-    unsafe {
-        while HBS_SPINLOCK.compare_and_swap(false, true, Ordering::SeqCst) {}
-    }
-    let result = work();
-    unsafe {
-        HBS_SPINLOCK.store(false, Ordering::SeqCst);
-    }
-    result
-}
 
 /// Initialize heartbeats
 pub fn init() {
+    lock_and_work(|hbs_opt|
+        if hbs_opt.is_none() {
             let mut hbs: Box<HashMap<ProfilerCategory, Heartbeat>> = Box::new(HashMap::new());
             maybe_create_heartbeat(&mut hbs, ProfilerCategory::ApplicationHeartbeat);
-    lock_and_work(||
+            *hbs_opt = Some(Box::into_raw(hbs))
-        unsafe {
-            HBS = Some(Box::into_raw(hbs));
         }
     );
 }
 
 /// Log regmaining buffer data and cleanup heartbeats
 pub fn cleanup() {
-    let hbs_opt: Option<Box<HashMap<ProfilerCategory, Heartbeat>>> = lock_and_work(||
+    let hbs_opt_box: Option<Box<HashMap<ProfilerCategory, Heartbeat>>> = lock_and_work(|hbs_opt|
+        hbs_opt.take().map(|hbs_ptr|
             unsafe {
-            match HBS {
+                Box::from_raw(hbs_ptr)
-                Some(hbs_ptr) => {
-                    let ret = Some(Box::from_raw(hbs_ptr));
-                    HBS = None;
-                    ret
-                },
-                None => None,
-            }
             }
+        )
     );
-    if let Some(mut hbs) = hbs_opt {
+    if let Some(mut hbs) = hbs_opt_box {
         for (_, mut v) in hbs.iter_mut() {
             // log any remaining heartbeat records before dropping
             log_heartbeat_records(v);
@@ -67,10 +43,12 @@ pub fn cleanup() {
 
 /// Check if a heartbeat exists for the given category
 pub fn is_heartbeat_enabled(category: &ProfilerCategory) -> bool {
-    let is_enabled = lock_and_work(||
+    let is_enabled = lock_and_work(|hbs_opt|
+        hbs_opt.map_or(false, |hbs_ptr|
             unsafe {
-            HBS.map_or(false, |hbs_ptr| (*hbs_ptr).contains_key(category))
+                (*hbs_ptr).contains_key(category)
             }
+        )
     );
     is_enabled || is_create_heartbeat(category)
 }
@@ -81,9 +59,9 @@ pub fn maybe_heartbeat(category: &ProfilerCategory,
                        end_time: u64,
                        start_energy: u64,
                        end_energy: u64) {
-    lock_and_work(||
+    lock_and_work(|hbs_opt|
+        if let Some(hbs_ptr) = *hbs_opt {
             unsafe {
-            if let Some(hbs_ptr) = HBS {
                 if !(*hbs_ptr).contains_key(category) {
                     maybe_create_heartbeat(&mut (*hbs_ptr), category.clone());
                 }
@@ -159,10 +137,33 @@ fn log_heartbeat_records(hb: &mut Heartbeat) {
     }
 }
 
+mod synchronized_heartbeat {
+    use heartbeats_simple::HeartbeatPow as Heartbeat;
+    use heartbeats_simple::HeartbeatPowContext as HeartbeatContext;
+    use profile_traits::time::ProfilerCategory;
+    use std::collections::HashMap;
+    use std::sync::atomic::{ATOMIC_BOOL_INIT, AtomicBool, Ordering};
+    use super::log_heartbeat_records;
+
+    static mut HBS: Option<*mut HashMap<ProfilerCategory, Heartbeat>> = None;
+
+    // unfortunately can't encompass the actual hashmap in a Mutex (Heartbeat isn't Send/Sync), so we'll use a spinlock
+    static HBS_SPINLOCK: AtomicBool = ATOMIC_BOOL_INIT;
+
+    pub fn lock_and_work<F, R>(work: F) -> R
+        where F: FnOnce(&mut Option<*mut HashMap<ProfilerCategory, Heartbeat>>) -> R {
+        while HBS_SPINLOCK.compare_and_swap(false, true, Ordering::SeqCst) {}
+        let result = unsafe {
+            work(&mut HBS)
+        };
+        HBS_SPINLOCK.store(false, Ordering::SeqCst);
+        result
+    }
+
     /// Callback function used to log the window buffer.
     /// When this is called from native C, the heartbeat is safely locked internally and the global lock is held.
     /// If calling from this file, you must already hold the global lock!
-extern fn heartbeat_window_callback(hb: *const HeartbeatContext) {
+    pub extern fn heartbeat_window_callback(hb: *const HeartbeatContext) {
         unsafe {
             if let Some(hbs_ptr) = HBS {
                 for (_, v) in (*hbs_ptr).iter_mut() {
@@ -173,3 +174,4 @@ extern fn heartbeat_window_callback(hb: *const HeartbeatContext) {
             }
         }
     }
+}