From 5589efe9bce8a9c38e941e7722ef07e681a2b491 Mon Sep 17 00:00:00 2001
From: porcuquine <porcuquine@users.noreply.github.com>
Date: Thu, 8 Oct 2020 10:05:36 -0700
Subject: [PATCH] Add fixes.

---
 .../porep/src/stacked/vanilla/cores.rs        | 35 +++++++++++--------
 .../src/stacked/vanilla/create_label/multi.rs | 10 +++---
 2 files changed, 26 insertions(+), 19 deletions(-)

diff --git a/storage-proofs/porep/src/stacked/vanilla/cores.rs b/storage-proofs/porep/src/stacked/vanilla/cores.rs
index 0df9b87c57..351ea34f9c 100644
--- a/storage-proofs/porep/src/stacked/vanilla/cores.rs
+++ b/storage-proofs/porep/src/stacked/vanilla/cores.rs
@@ -42,7 +42,26 @@ pub fn checkout_core_group() -> Option<MutexGuard<'static, CoreGroup>> {
     }
 }
 
+#[cfg(not(target_os = "windows"))]
+pub type ThreadId = libc::pthread_t;
+
+#[cfg(target_os = "windows")]
+pub type ThreadId = winapi::winnt::HANDLE;
+
+/// Helper method to get the thread id through libc, with current rust stable (1.5.0) its not
+/// possible otherwise I think.
+#[cfg(not(target_os = "windows"))]
+fn get_thread_id() -> ThreadId {
+    unsafe { libc::pthread_self() }
+}
+
+#[cfg(target_os = "windows")]
+fn get_thread_id() -> ThreadId {
+    unsafe { kernel32::GetCurrentThread() }
+}
+
 pub struct Cleanup {
+    tid: ThreadId,
     prior_state: Option<hwloc::Bitmap>,
 }
 
@@ -51,9 +70,8 @@ impl Drop for Cleanup {
         match self.prior_state.take() {
             Some(prior) => {
                 let child_topo = &TOPOLOGY;
-                let tid = get_thread_id();
                 let mut locked_topo = child_topo.lock().unwrap();
-                let _ = locked_topo.set_cpubind_for_thread(tid, prior, CPUBIND_THREAD);
+                let _ = locked_topo.set_cpubind_for_thread(self.tid, prior, CPUBIND_THREAD);
             }
             None => (),
         }
@@ -91,6 +109,7 @@ pub fn bind_core(core_index: CoreIndex) -> Result<Cleanup> {
     }
 
     Ok(Cleanup {
+        tid,
         prior_state: before,
     })
 }
@@ -109,18 +128,6 @@ fn get_core_by_index<'a>(topo: &'a Topology, index: CoreIndex) -> Result<&'a Top
     }
 }
 
-/// Helper method to get the thread id through libc, with current rust stable (1.5.0) its not
-/// possible otherwise I think.
-#[cfg(not(target_os = "windows"))]
-fn get_thread_id() -> libc::pthread_t {
-    unsafe { libc::pthread_self() }
-}
-
-#[cfg(target_os = "windows")]
-fn get_thread_id() -> winapi::winnt::HANDLE {
-    unsafe { kernel32::GetCurrentThread() }
-}
-
 fn core_groups(cores_per_unit: usize) -> Option<Vec<Mutex<Vec<CoreIndex>>>> {
     let topo = TOPOLOGY.lock().unwrap();
 
diff --git a/storage-proofs/porep/src/stacked/vanilla/create_label/multi.rs b/storage-proofs/porep/src/stacked/vanilla/create_label/multi.rs
index d234de858c..4606601aad 100644
--- a/storage-proofs/porep/src/stacked/vanilla/create_label/multi.rs
+++ b/storage-proofs/porep/src/stacked/vanilla/create_label/multi.rs
@@ -541,13 +541,13 @@ pub fn create_labels_for_decoding<Tree: 'static + MerkleTreeTrait, T: AsRef<[u8]
 
     let core_group = Arc::new(checkout_core_group());
 
-    if let Some(group) = &*core_group {
+    // When `_cleanup_handle` is dropped, the previous binding of thread will be restored.
+    let _cleanup_handle = (*core_group).as_ref().map(|group| {
         // This could fail, but we will ignore the error if so.
         // It will be logged as a warning by `bind_core`.
-        if let Some(&core_index) = group.get(0) {
-            let _ = bind_core(core_index);
-        }
-    };
+        debug!("binding core in main thread");
+        group.get(0).map(|core_index| bind_core(*core_index))
+    });
 
     // NOTE: this means we currently keep 2x sector size around, to improve speed
     let (parents_cache, mut layer_labels, mut exp_labels) = setup_create_label_memory(