chore: Enforce usage of parking_lot (#3325)

Mostly because of better ergonomics, but also has better performance in
some cases and some useful additional features.

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

.cargo/config.toml

@@ -7,7 +7,7 @@ rustflags = ["-C", "link-arg=-undefined", "-C", "link-arg=dynamic_lookup"]
 [target.x86_64-apple-darwin]
 rustflags = ["-C", "link-arg=-undefined", "-C", "link-arg=dynamic_lookup"]
 [target.wasm32-unknown-unknown]
-rustflags = ['--cfg', 'getrandom_backend="wasm_js"']
+rustflags = ['--cfg', 'getrandom_backend="wasm_js"', '-C', 'target-feature=+atomics']
 
 [alias]
 xtask = "run -p xtask --"

.github/workflows/ci.yml

@@ -419,7 +419,7 @@ jobs:
           targets: "wasm32-wasip1"
       - name: Setup Wasmer
         uses: wasmerio/[email protected]
-      - run: cargo build --target wasm32-wasip1
+      - run: cargo -Zbuild-std=panic_abort,std build --target wasm32-wasip1
         working-directory: ./wasm-test
       - run: wasmer run ./target/wasm32-wasip1/debug/wasm-test.wasm
         working-directory: ./wasm-test

Cargo.lock

@@ -122,7 +122,6 @@ object_store = "0.12.1"
 opentelemetry = "0.29.0"
 opentelemetry-otlp = "0.29.0"
 opentelemetry_sdk = "0.29.0"
-parking_lot = "0.12.3"
 parquet = "55"
 paste = "1.0.15"
 pin-project = "1.1.5"
@@ -134,6 +133,7 @@ pyo3 = { version = "0.24.1", features = ["extension-module", "abi3-py310"] }
 pyo3-log = "0.12.1"
 rancor = "0.1.0"
 rand = "0.9.0"
+parking_lot = { version = "0.12.3", features = ["nightly"] }
 rand_distr = "0.5"
 ratatui = "0.29"
 rayon = "1.10.0"

Cargo.toml

@@ -2,4 +2,5 @@ allow-dbg-in-tests = true
 allow-expect-in-tests = true
 allow-unwrap-in-tests = true
 single-char-binding-names-threshold = 2
-disallowed-types = ["std::collections::HashMap", "std::collections::HashSet"]
+# We prefer using parking_lot for improved ergonomics and performance.
+disallowed-types = ["std::collections::HashMap", "std::collections::HashSet", "std::sync::Mutex", "std::sync::RwLock"]

clippy.toml

@@ -29,6 +29,7 @@ itertools = { workspace = true }
 log = { workspace = true }
 mimalloc = { workspace = true }
 object_store = { workspace = true, features = ["aws", "gcp", "azure", "http"] }
+parking_lot = { workspace = true }
 pyo3 = { workspace = true, features = ["abi3-py311"] }
 pyo3-log = { workspace = true }
 tokio = { workspace = true, features = ["fs", "rt-multi-thread"] }

pyvortex/Cargo.toml

@@ -2,16 +2,15 @@ mod python;
 pub(crate) mod stream;
 
 use std::iter;
-use std::sync::Mutex;
 
 use arrow::array::RecordBatchReader;
 use arrow::pyarrow::IntoPyArrow;
+use parking_lot::Mutex;
 use pyo3::prelude::*;
 use pyo3::types::PyIterator;
 use pyo3::{Bound, PyResult, Python};
 pub(crate) use stream::*;
 use vortex::dtype::DType;
-use vortex::error::VortexExpect;
 use vortex::iter::{ArrayIterator, ArrayIteratorAdapter, ArrayIteratorExt};
 use vortex::{Canonical, IntoArray};
 
@@ -51,7 +50,7 @@ impl PyArrayIterator {
     }
 
     pub fn take(&self) -> Option<Box<dyn ArrayIterator + Send>> {
-        self.iter.lock().vortex_expect("poisoned lock").take()
+        self.iter.lock().take()
     }
 }
 
@@ -75,7 +74,6 @@ impl PyArrayIterator {
             Ok(self
                 .iter
                 .lock()
-                .vortex_expect("poisoned lock")
                 .as_mut()
                 .and_then(|iter| iter.next())
                 .transpose()?
@@ -87,7 +85,7 @@ impl PyArrayIterator {
     /// this will be a :class:`vortex.ChunkedArray`, otherwise it will be a single array.
     fn read_all(&self, py: Python) -> PyResult<PyArrayRef> {
         let array = py.allow_threads(|| {
-            if let Some(iter) = self.iter.lock().vortex_expect("poisoned lock").take() {
+            if let Some(iter) = self.iter.lock().take() {
                 iter.read_all()
             } else {
                 // Otherwise, we continue to return an empty array.

pyvortex/src/iter/mod.rs

@@ -1,10 +1,10 @@
-use std::sync::{Arc, RwLock};
+use std::sync::Arc;
 
 use itertools::Itertools;
+use parking_lot::RwLock;
 use pyo3::prelude::*;
 use pyo3::{Bound, PyResult, Python};
 use vortex::ArrayRegistry;
-use vortex::error::VortexExpect;
 use vortex::file::DEFAULT_REGISTRY;
 use vortex::layout::{LayoutRegistry, LayoutRegistryExt};
 
@@ -58,16 +58,13 @@ impl PyRegistry {
     /// It's not currently possible to register a layout encoding from Python.
     pub(crate) fn register(&self, cls: PythonEncoding) -> PyResult<()> {
         let encoding = cls.to_encoding();
-        self.array_registry
-            .write()
-            .vortex_expect("poisoned lock")
-            .register(encoding);
+        self.array_registry.write().register(encoding);
         Ok(())
     }
 
     /// Create an :class:`~vortex.ArrayContext` containing the given encodings.
     fn array_ctx(&self, encodings: Vec<Bound<PyAny>>) -> PyResult<PyArrayContext> {
-        let registry = self.array_registry.read().vortex_expect("poisoned lock");
+        let registry = self.array_registry.read();
         Ok(PyArrayContext::from(registry.new_context(
             encoding_ids(&encodings)?.iter().map(|s| s.as_str()),
         )?))

pyvortex/src/registry.rs

@@ -8,7 +8,6 @@
 
 use std::any::{Any, type_name};
 use std::fmt::{Debug, Formatter};
-use std::sync::RwLock;
 
 use arcref::ArcRef;
 pub use between::*;
@@ -27,10 +26,11 @@ pub use mask::*;
 pub use min_max::*;
 pub use nan_count::*;
 pub use numeric::*;
+use parking_lot::RwLock;
 pub use sum::*;
 pub use take::*;
 use vortex_dtype::DType;
-use vortex_error::{VortexError, VortexExpect, VortexResult, vortex_bail, vortex_err};
+use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
 use vortex_mask::Mask;
 use vortex_scalar::Scalar;
 
@@ -84,10 +84,7 @@ impl ComputeFn {
 
     /// Register a kernel for the compute function.
     pub fn register_kernel(&self, kernel: ArcRef<dyn Kernel>) {
-        self.kernels
-            .write()
-            .vortex_expect("poisoned lock")
-            .push(kernel);
+        self.kernels.write().push(kernel);
     }
 
     /// Invokes the compute function with the given arguments.
@@ -112,9 +109,7 @@ impl ComputeFn {
         let expected_dtype = self.vtable.return_dtype(args)?;
         let expected_len = self.vtable.return_len(args)?;
 
-        let output = self
-            .vtable
-            .invoke(args, &self.kernels.read().vortex_expect("poisoned lock"))?;
+        let output = self.vtable.invoke(args, &self.kernels.read())?;
 
         if output.dtype() != &expected_dtype {
             vortex_bail!(
@@ -154,7 +149,7 @@ impl ComputeFn {
 
     /// Returns the compute function's kernels.
     pub fn kernels(&self) -> Vec<ArcRef<dyn Kernel>> {
-        self.kernels.read().vortex_expect("poisoned lock").to_vec()
+        self.kernels.read().to_vec()
     }
 }
 

vortex-array/src/compute/mod.rs

@@ -1,7 +1,8 @@
 use std::fmt::Display;
-use std::sync::{Arc, RwLock, RwLockReadGuard};
+use std::sync::Arc;
 
 use itertools::Itertools;
+use parking_lot::RwLock;
 use vortex_error::{VortexExpect, VortexResult, vortex_err};
 
 use crate::EncodingRef;
@@ -56,7 +57,7 @@ impl<T: Clone + Eq> VTableContext<T> {
 
     pub fn with(self, encoding: T) -> Self {
         {
-            let mut write = self.0.write().vortex_expect("poisoned lock");
+            let mut write = self.0.write();
             if write.iter().all(|e| e != &encoding) {
                 write.push(encoding);
             }
@@ -68,13 +69,13 @@ impl<T: Clone + Eq> VTableContext<T> {
         items.into_iter().fold(self, |ctx, e| ctx.with(e))
     }
 
-    pub fn encodings(&self) -> RwLockReadGuard<Vec<T>> {
-        self.0.read().vortex_expect("poisoned lock")
+    pub fn encodings(&self) -> Vec<T> {
+        self.0.read().clone()
     }
 
     /// Returns the index of the encoding in the context, or adds it if it doesn't exist.
     pub fn encoding_idx(&self, encoding: &T) -> u16 {
-        let mut write = self.0.write().vortex_expect("poisoned lock");
+        let mut write = self.0.write();
         if let Some(idx) = write.iter().position(|e| e == encoding) {
             return u16::try_from(idx).vortex_expect("Cannot have more than u16::MAX encodings");
         }
@@ -88,11 +89,7 @@ impl<T: Clone + Eq> VTableContext<T> {
 
     /// Find an encoding by its position.
     pub fn lookup_encoding(&self, idx: u16) -> Option<T> {
-        self.0
-            .read()
-            .vortex_expect("poisoned lock")
-            .get(idx as usize)
-            .cloned()
+        self.0.read().get(idx as usize).cloned()
     }
 }