use global mutex for gil

Author: wangxiaoying

Cargo.lock

@@ -1,3 +1,5 @@
+set dotenv-load := true
+
 build-release:
     cargo build  --release
 

Justfile

@@ -37,6 +37,7 @@ tokio = {version = "1", features = ["rt", "rt-multi-thread", "net"]}
 tokio-util = "0.6"
 url = "2"
 uuid = "0.8"
+lazy_static = "1.4.0"
 
 [build-dependencies]
 built = {version = "0.5", features = ["chrono"]}

connectorx-python/Cargo.toml

@@ -11,6 +11,9 @@ use pyo3::prelude::*;
 use pyo3::{wrap_pyfunction, PyResult};
 use std::sync::Once;
 
+#[macro_use]
+extern crate lazy_static;
+
 static START: Once = Once::new();
 
 // https://github.com/PyO3/pyo3-built/issues/21

connectorx-python/src/lib.rs

@@ -1,4 +1,4 @@
-use super::{check_dtype, HasPandasColumn, PandasColumn, PandasColumnObject};
+use super::{check_dtype, HasPandasColumn, PandasColumn, PandasColumnObject, GIL_MUTEX};
 use crate::errors::ConnectorXPythonError;
 use anyhow::anyhow;
 use fehler::throws;
@@ -7,7 +7,6 @@ use numpy::{npyffi::NPY_TYPES, Element, PyArray, PyArrayDescr};
 use pyo3::{FromPyObject, Py, PyAny, PyResult, Python, ToPyObject};
 use std::any::TypeId;
 use std::marker::PhantomData;
-use std::sync::{Arc, Mutex};
 
 #[derive(Clone)]
 #[repr(transparent)]
@@ -23,7 +22,6 @@ impl Element for PyList {
 
 pub struct ArrayBlock<'a, V> {
     data: ArrayViewMut2<'a, PyList>,
-    mutex: Arc<Mutex<()>>,
     buf_size_mb: usize,
     _value_type: PhantomData<V>,
 }
@@ -35,8 +33,7 @@ impl<'a, V> FromPyObject<'a> for ArrayBlock<'a, V> {
         let data = unsafe { array.as_array_mut() };
         Ok(ArrayBlock::<V> {
             data,
-            mutex: Arc::new(Mutex::new(())), // allocate the lock here since only a few blocks needs to aquire the GIL for now
-            buf_size_mb: 16,                 // in MB
+            buf_size_mb: 16, // in MB
             _value_type: PhantomData,
         })
     }
@@ -61,7 +58,6 @@ impl<'a, V> ArrayBlock<'a, V> {
                 lengths: vec![],
                 buffer: Vec::with_capacity(self.buf_size_mb * (1 << 17) * 11 / 10), // allocate a little bit more memory to avoid Vec growth
                 buf_size: self.buf_size_mb * (1 << 17),
-                mutex: self.mutex.clone(),
             })
         }
         ret
@@ -74,7 +70,6 @@ pub struct ArrayColumn<'a, V> {
     buffer: Vec<V>,
     lengths: Vec<usize>, // usize::MAX if the string is None
     buf_size: usize,
-    mutex: Arc<Mutex<()>>,
 }
 
 impl<'a, V> PandasColumnObject for ArrayColumn<'a, V>
@@ -179,7 +174,6 @@ where
                 lengths: vec![],
                 buffer: Vec::with_capacity(self.buf_size),
                 buf_size: self.buf_size,
-                mutex: self.mutex.clone(),
             });
         }
         partitions
@@ -194,8 +188,7 @@ where
 
             {
                 // allocation in python is not thread safe
-                let _guard = self
-                    .mutex
+                let _guard = GIL_MUTEX
                     .lock()
                     .map_err(|e| anyhow!("mutex poisoned {}", e))?;
                 let mut start = 0;

connectorx-python/src/pandas/pandas_columns/array.rs

@@ -1,12 +1,11 @@
-use super::{check_dtype, HasPandasColumn, PandasColumn, PandasColumnObject};
+use super::{check_dtype, HasPandasColumn, PandasColumn, PandasColumnObject, GIL_MUTEX};
 use crate::errors::ConnectorXPythonError;
 use anyhow::anyhow;
 use fehler::throws;
 use ndarray::{ArrayViewMut2, Axis, Ix2};
 use numpy::{npyffi::NPY_TYPES, Element, PyArray, PyArrayDescr};
 use pyo3::{FromPyObject, Py, PyAny, PyResult, Python};
 use std::any::TypeId;
-use std::sync::{Arc, Mutex};
 
 #[derive(Clone)]
 #[repr(transparent)]
@@ -22,7 +21,6 @@ impl Element for PyBytes {
 
 pub struct BytesBlock<'a> {
     data: ArrayViewMut2<'a, PyBytes>,
-    mutex: Arc<Mutex<()>>,
     buf_size_mb: usize,
 }
 
@@ -33,8 +31,7 @@ impl<'a> FromPyObject<'a> for BytesBlock<'a> {
         let data = unsafe { array.as_array_mut() };
         Ok(BytesBlock {
             data,
-            mutex: Arc::new(Mutex::new(())), // allocate the lock here since only a few blocks needs to aquire the GIL for now
-            buf_size_mb: 16,                 // in MB
+            buf_size_mb: 16, // in MB
         })
     }
 }
@@ -58,7 +55,6 @@ impl<'a> BytesBlock<'a> {
                 bytes_lengths: vec![],
                 bytes_buf: Vec::with_capacity(self.buf_size_mb * (1 << 20) * 11 / 10), // allocate a little bit more memory to avoid Vec growth
                 buf_size: self.buf_size_mb * (1 << 20),
-                mutex: self.mutex.clone(),
             })
         }
         ret
@@ -71,7 +67,6 @@ pub struct BytesColumn<'a> {
     bytes_buf: Vec<u8>,
     bytes_lengths: Vec<usize>, // usize::MAX if the string is None
     buf_size: usize,
-    mutex: Arc<Mutex<()>>,
 }
 
 impl<'a> PandasColumnObject for BytesColumn<'a> {
@@ -171,7 +166,6 @@ impl<'a> BytesColumn<'a> {
                 bytes_lengths: vec![],
                 bytes_buf: Vec::with_capacity(self.buf_size),
                 buf_size: self.buf_size,
-                mutex: self.mutex.clone(),
             });
         }
 
@@ -187,8 +181,7 @@ impl<'a> BytesColumn<'a> {
 
             {
                 // allocation in python is not thread safe
-                let _guard = self
-                    .mutex
+                let _guard = GIL_MUTEX
                     .lock()
                     .map_err(|e| anyhow!("mutex poisoned {}", e))?;
                 let mut start = 0;

connectorx-python/src/pandas/pandas_columns/bytes.rs

@@ -17,8 +17,14 @@ pub use float64::{Float64Block, Float64Column};
 pub use int64::{Int64Block, Int64Column};
 use pyo3::{exceptions::PyRuntimeError, PyAny, PyResult};
 use std::any::TypeId;
+use std::sync::Mutex;
 pub use string::{StringBlock, StringColumn};
 
+// A global GIL lock for Python object allocations like string, bytes and list
+lazy_static! {
+    static ref GIL_MUTEX: Mutex<()> = Mutex::new(());
+}
+
 pub trait PandasColumnObject: Send {
     fn typecheck(&self, _: TypeId) -> bool;
     fn typename(&self) -> &'static str;

connectorx-python/src/pandas/pandas_columns/mod.rs

@@ -1,5 +1,5 @@
 use super::super::pystring::{PyString, StringInfo};
-use super::{check_dtype, HasPandasColumn, PandasColumn, PandasColumnObject};
+use super::{check_dtype, HasPandasColumn, PandasColumn, PandasColumnObject, GIL_MUTEX};
 use crate::errors::ConnectorXPythonError;
 use anyhow::anyhow;
 use fehler::throws;
@@ -8,11 +8,9 @@ use ndarray::{ArrayViewMut2, Axis, Ix2};
 use numpy::PyArray;
 use pyo3::{FromPyObject, PyAny, PyResult, Python};
 use std::any::TypeId;
-use std::sync::{Arc, Mutex};
 
 pub struct StringBlock<'a> {
     data: ArrayViewMut2<'a, PyString>,
-    mutex: Arc<Mutex<()>>,
     buf_size_mb: usize,
 }
 
@@ -23,8 +21,7 @@ impl<'a> FromPyObject<'a> for StringBlock<'a> {
         let data = unsafe { array.as_array_mut() };
         Ok(StringBlock {
             data,
-            mutex: Arc::new(Mutex::new(())), // allocate the lock here since only a few blocks needs to aquire the GIL for now
-            buf_size_mb: 4,                  // in MB
+            buf_size_mb: 4, // in MB
         })
     }
 }
@@ -48,7 +45,6 @@ impl<'a> StringBlock<'a> {
                 string_lengths: vec![],
                 string_buf: Vec::with_capacity(self.buf_size_mb * (1 << 20) * 11 / 10), // allocate a little bit more memory to avoid Vec growth
                 buf_size: self.buf_size_mb * (1 << 20),
-                mutex: self.mutex.clone(),
             })
         }
         ret
@@ -61,7 +57,6 @@ pub struct StringColumn<'a> {
     string_buf: Vec<u8>,
     string_lengths: Vec<usize>, // usize::MAX for empty string
     buf_size: usize,
-    mutex: Arc<Mutex<()>>,
 }
 
 impl<'a> PandasColumnObject for StringColumn<'a> {
@@ -236,7 +231,6 @@ impl<'a> StringColumn<'a> {
                 string_lengths: vec![],
                 string_buf: Vec::with_capacity(self.buf_size),
                 buf_size: self.buf_size,
-                mutex: self.mutex.clone(),
             });
         }
 
@@ -252,11 +246,11 @@ impl<'a> StringColumn<'a> {
 
         let py = unsafe { Python::assume_gil_acquired() };
         let _guard = if force {
-            self.mutex
+            GIL_MUTEX
                 .lock()
                 .map_err(|e| anyhow!("mutex poisoned {}", e))?
         } else {
-            match self.mutex.try_lock() {
+            match GIL_MUTEX.try_lock() {
                 Ok(guard) => guard,
                 Err(_) => return,
             }