core: Deal with drop(ThreadPool) from a local operation.

From other injected jobs this is not an issue, we'll end up waiting for
ever till out of work.

rayon-core/src/registry.rs

@@ -518,10 +518,7 @@ impl Registry {
             } else if (*worker_thread).registry().id() != self.id() {
                 self.in_worker_cross(&*worker_thread, op)
             } else {
-                // Perfectly valid to give them a `&T`: this is the
-                // current thread, so we know the data structure won't be
-                // invalidated until we return.
-                op(&*worker_thread, false)
+                (*worker_thread).execute_local_op(op)
             }
         }
     }
@@ -630,6 +627,9 @@ impl Registry {
         if t.index != 0 {
             return;
         }
+        if t.in_local_op.get() {
+            return;
+        }
         unsafe {
             wait_until_out_of_work(t);
             let _ = Box::from_raw(WorkerThread::current() as *mut WorkerThread);
@@ -698,6 +698,9 @@ pub(super) struct WorkerThread {
     rng: XorShift64Star,
 
     registry: Arc<Registry>,
+
+    /// Whether we're executing a local operation that must keep us alive.
+    in_local_op: Cell<bool>,
 }
 
 // This is a bit sketchy, but basically: the WorkerThread is
@@ -718,6 +721,7 @@ impl From<ThreadBuilder> for WorkerThread {
             index: thread.index,
             rng: XorShift64Star::new(),
             registry: thread.registry,
+            in_local_op: Cell::new(false),
         }
     }
 }
@@ -946,6 +950,21 @@ impl WorkerThread {
             }
         }
     }
+
+    fn execute_local_op<OP, R>(&self, op: OP) -> R
+    where
+        OP: FnOnce(&WorkerThread, bool) -> R + Send,
+        R: Send,
+    {
+        let old_in_local_op = self.in_local_op.get();
+        self.in_local_op.set(true);
+
+        let r = op(self, false);
+
+        self.in_local_op.set(old_in_local_op);
+
+        r
+    }
 }
 
 /// ////////////////////////////////////////////////////////////////////////
@@ -1017,10 +1036,7 @@ where
     unsafe {
         let owner_thread = WorkerThread::current();
         if !owner_thread.is_null() {
-            // Perfectly valid to give them a `&T`: this is the
-            // current thread, so we know the data structure won't be
-            // invalidated until we return.
-            op(&*owner_thread, false)
+            (*owner_thread).execute_local_op(op)
         } else {
             global_registry().in_worker(op)
         }

rayon-core/tests/use_current_thread.rs

@@ -46,3 +46,20 @@ fn use_current_thread_basic() {
         let _ = handle.join();
     }
 }
+
+#[test]
+fn use_current_thread_drop_from_pool() {
+    let pool = ThreadPoolBuilder::new()
+        .use_current_thread()
+        .build()
+        .unwrap();
+    assert_eq!(rayon_core::current_thread_index(), Some(0));
+
+    rayon_core::scope(|_| drop(pool));
+
+    assert_eq!(
+        rayon_core::current_thread_index(),
+        Some(0),
+        "Can't shut down current thread cleanly if in a local op"
+    );
+}