Fix hanging core-relations tests. (#28)

* Fix deadlock condition and add mimalloc

* remove a different deadlock source

---------

Co-authored-by: Eli Rosenthal <[email protected]>

Author: ezrosent

core-relations/Cargo.toml

@@ -31,6 +31,7 @@ dyn-clone = "1.0.17"
 [dev-dependencies]
 once_cell = "1.20.2"
 divan = "0.1.15"
+mimalloc = "0.1"
 
 [[bench]]
 name = "writes"

core-relations/src/free_join/mod.rs

@@ -647,36 +647,44 @@ impl Drop for Database {
 /// This is in a separate function to allow us to reuse it while already
 /// borrowing a `TableInfo`.
 fn get_index_from_tableinfo(table_info: &TableInfo, cols: &[ColumnId]) -> HashIndex {
-    let guard = table_info.indexes.entry(cols.into()).or_insert_with(|| {
-        Arc::new(ReadOptimizedLock::new(Index::new(
-            cols.to_vec(),
-            TupleIndex::new(cols.len()),
-        )))
-    });
-    let ix = guard.value().read();
+    let index: Arc<_> = table_info
+        .indexes
+        .entry(cols.into())
+        .or_insert_with(|| {
+            Arc::new(ReadOptimizedLock::new(Index::new(
+                cols.to_vec(),
+                TupleIndex::new(cols.len()),
+            )))
+        })
+        .clone();
+    let ix = index.read();
     if ix.needs_refresh(table_info.table.as_ref()) {
         mem::drop(ix);
-        let mut ix = guard.value().lock();
+        let mut ix = index.lock();
         ix.refresh(table_info.table.as_ref());
     }
-    guard.value().clone()
+    index
 }
 
 /// The core logic behind getting and updating a column index.
 ///
 /// This is the single-column analog to [`get_index_from_tableinfo`].
 fn get_column_index_from_tableinfo(table_info: &TableInfo, col: ColumnId) -> HashColumnIndex {
-    let index = table_info.column_indexes.entry(col).or_insert_with(|| {
-        Arc::new(ReadOptimizedLock::new(Index::new(
-            vec![col],
-            ColumnIndex::new(),
-        )))
-    });
+    let index: Arc<_> = table_info
+        .column_indexes
+        .entry(col)
+        .or_insert_with(|| {
+            Arc::new(ReadOptimizedLock::new(Index::new(
+                vec![col],
+                ColumnIndex::new(),
+            )))
+        })
+        .clone();
     let ix = index.read();
     if ix.needs_refresh(table_info.table.as_ref()) {
         mem::drop(ix);
         let mut ix = index.lock();
         ix.refresh(table_info.table.as_ref());
     }
-    index.clone()
+    index
 }

core-relations/src/hash_index/mod.rs

@@ -2,9 +2,10 @@
 use std::{
     hash::{Hash, Hasher},
     mem,
-    sync::Mutex,
+    sync::{Arc, Mutex},
 };
 
+use concurrency::Notification;
 use hashbrown::HashTable;
 use numeric_id::{define_id, IdVec, NumericId};
 use once_cell::sync::Lazy;
@@ -260,22 +261,23 @@ impl IndexBase for ColumnIndex {
             }
         };
 
-        THREAD_POOL.install(|| {
-            rayon::scope(|scope| {
+        run_in_thread_pool_and_block(&THREAD_POOL, || {
+            rayon::in_place_scope(|inner| {
                 let mut cur = Offset::new(0);
                 loop {
                     let mut buf = TaggedRowBuffer::new(cols.len());
                     if let Some(next) =
                         table.scan_project(subset, cols, cur, BATCH_SIZE, &[], &mut buf)
                     {
                         cur = next;
-                        scope.spawn(move |_| split_buf(buf));
+                        inner.spawn(move |_| split_buf(buf));
                     } else {
-                        scope.spawn(move |_| split_buf(buf));
+                        inner.spawn(move |_| split_buf(buf));
                         break;
                     }
                 }
             });
+
             self.shards.par_iter_mut().for_each(|(shard_id, shard)| {
                 use indexmap::map::Entry;
                 // Sort the vector by start row id to ensure we populate subsets in sorted order.
@@ -302,6 +304,43 @@ impl IndexBase for ColumnIndex {
     }
 }
 
+/// This function is an alternative for [`rayon::ThreadPool::install`] that doesn't steal work from
+/// the callee's current thread pool while waiting for `f` to finish.
+///
+/// We do this to avoid deadlocks. The whole purpose of using a separate threadpool in this module
+/// is to allow for sufficient parallelism while holding a lock on the main threadpool. That means
+/// we are not worried about an outer lock tying up a thread in the main pool.
+///
+/// On the other hand, it _is_ a bad idea to steal work on a rayon thread pool with some locks
+/// held. In particular, if another task on the thread pool _itself_ attempts to aquire the same
+/// lock, this can cause a deadlock. We saw this in the tests for this crate. The relevant lock
+/// are those around individual indexes stored in the database-level index cache.
+fn run_in_thread_pool_and_block<'a>(pool: &rayon::ThreadPool, f: impl FnMut() + Send + 'a) {
+    // NB: We don't need the heap allocations here. But we are only calling this function if
+    // we are about to do a bunch of work, so clarify is probably going to be better than (even
+    // more) unsafe code.
+
+    // Alright, here we go: pretend `f` has `'static` lifetime because we are passing it to
+    // `spawn`.
+    trait LifetimeWork<'a>: FnMut() + Send + 'a {}
+
+    impl<'a, F: FnMut() + Send + 'a> LifetimeWork<'a> for F {}
+    let as_lifetime: Box<dyn LifetimeWork<'a>> = Box::new(f);
+    let mut casted_away = unsafe {
+        // SAFETY: `casted_away` will be dropped at the end of this method. The notification used
+        // below will ensure it does not escape.
+        mem::transmute::<Box<dyn LifetimeWork<'a>>, Box<dyn LifetimeWork<'static>>>(as_lifetime)
+    };
+    let n = Arc::new(Notification::new());
+    let inner = n.clone();
+    pool.spawn(move || {
+        casted_away();
+        mem::drop(casted_away);
+        inner.notify();
+    });
+    n.wait()
+}
+
 impl ColumnIndex {
     pub(crate) fn new() -> ColumnIndex {
         with_pool_set(|ps| {
@@ -447,7 +486,7 @@ impl IndexBase for TupleIndex {
                 queues[shard_id].lock().unwrap().push((first, buf));
             }
         };
-        THREAD_POOL.install(|| {
+        run_in_thread_pool_and_block(&THREAD_POOL, || {
             rayon::scope(|scope| {
                 let mut cur = Offset::new(0);
                 loop {

core-relations/src/tests.rs

@@ -19,6 +19,11 @@ use crate::{
     PlanStrategy,
 };
 
+/// On MacOs the system allocator is vulenrable to contention, causing tests to execute quite
+/// slowly without mimalloc.
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
 #[test]
 fn basic_query() {
     let MathEgraph {