Thread Safe Decorator <Help Wanted>

[zachcoleman]

Building upon the example in the repo, I tried calling the decorator from a ThreadPoolExecutor and run into Runtime Error: Already Borrowed.

Here's the Rust:

use pyo3::prelude::*;
use pyo3::types::{PyDict, PyTuple};

#[pyclass(name = "Counter")]
pub struct Counter{
    #[pyo3(get)]
    count: usize,
    wraps: Py<PyAny>,
}

#[pymethods]
impl Counter{
    #[new]
    fn __new__(wraps: Py<PyAny>) -> Self {
        Counter { 
            count: 0,
            wraps: wraps,
        }
    }

    #[args(args = "*", kwargs = "**")]
    fn __call__(
        &mut self,
        py: Python<'_>,
        args: &PyTuple,
        kwargs: Option<&PyDict>,
    ) -> PyResult<Py<PyAny>> {
        self.count += 1;
        return self.wraps.call(py, args, kwargs);
    }
}

#[pymodule]
pub fn rust_counter(_py: Python<'_>, module: &PyModule) -> PyResult<()> {
    module.add_class::<Counter>()?;
    Ok(())
}

Calling from ThreadPoolExecutor:

from concurrent.futures import ThreadPoolExecutor
from rust_counter import Counter 

@Counter
def hello():
    print("hello")

with ThreadPoolExecutor(4) as ex:
    futures = []
    for _ in range(4):
        futures.append(ex.submit(hello))

Results in me getting: RuntimeError: Already borrowed

I've tried a few methods:

• explicitly using Python::with_gil in the __call__ method
• adding a Rust mutex to the count variable (including an Arc<Mutex<usize>> to it)
• adding py.allow_threads around the self.count += 1 call (not sure why this would've helped, but I tried)

I was just wondering if anyone knows how to make this kind of pattern thread safe. I thought the __call__ method accepting the Python token as an arg could serve as proof of holding GIL, but somehow I'm running into an unsafe operation. If anyone has any insight, it would be much appreciated.

Based on similar issues, this may be a problem w/ the &mut self pattern used here. If that's the case, has anyone figured out a way around this... is it a case of the GIL not preventing this method from being called simultaneously (i.e. __call__ is not atomic and so two calls to this method are happening at the same time)?

[zachcoleman]

Okay I will leave this question up for posteriority, but I was able to get a solution working.

Rust:

use pyo3::prelude::*;
use pyo3::types::{PyDict, PyTuple};
use std::sync::{Arc, Mutex};

#[pyclass(name = "Counter")]
pub struct Counter{
    // #[pyo3(get)]
    count: Arc<Mutex<usize>>,
    wraps: Py<PyAny>,
}

#[pymethods]
impl Counter{
    #[new]
    fn __new__(wraps: Py<PyAny>) -> Self {
        Counter { 
            count: Arc::new(Mutex::new(0)),
            wraps: wraps,
        }
    }

    fn count(&self) -> usize{
        return *self.count.lock().unwrap()
    }

    #[args(args = "*", kwargs = "**")]
    fn __call__(
        &self,
        py: Python<'_>,
        args: &PyTuple,
        kwargs: Option<&PyDict>,
    ) -> PyResult<Py<PyAny>> {
        {
            let mut count = self.count.lock().unwrap();
            *count += 1;
        }
        return self.wraps.call(py, args, kwargs);
    }
}

#[pymodule]
pub fn rust_counter(_py: Python<'_>, module: &PyModule) -> PyResult<()> {
    module.add_class::<Counter>()?;
    Ok(())
}

I'm unsure why, but the additional scope around the count incrementing is crucial. If not, the process deadlocks and I believe this comes from the count not going out of scope and dropping its lock too late (?). Also by using an interior mutability pattern in this method, the __call__ method can now accept &self and does not need &mut self.

Caveats:

• Arc + Mutex adds more overhead vs Rc + RefCell so, this only should be used for threaded environments
• Arc<Mutex<T>> can not use the macro #[pyo3(get)], so a quick getter method was added count()

Testing code:

from concurrent.futures import ThreadPoolExecutor
from rust_counter import Counter 

@Counter
def hello():
    return "hello"

num_threads, num_runs = 10, 100_000
with ThreadPoolExecutor(num_threads) as ex:
    futures = []
    for _ in range(num_runs):
        futures.append(ex.submit(hello))

assert len([f for f in futures if f.exception()]) == 0
assert hello.count() == num_runs

Feedback welcome!

[davidhewitt]

Two comments:

1. My experience would tell me the deadlock occurs because the Python GIL is effectively another mutex. The .lock() in __call__ will block the Python thread meaning that thread doesn't release the GIL, and so another thread holding the Mutex lock will deadlock on attempting to acquire the GIL. The GIL release/acquire cycle can happen inside self.wraps.call(py, args, kwargs) inside the Python interpreter to try and give other Python threads a chance to run. This would explain why your explicit scope to drop the lock before calling self.wraps fixes the issue. In general, if using Mutex, I suggest releasing the GIL before locking with py.allow_threads. E.g.

    #[args(args = "*", kwargs = "**")]
    fn __call__(
        &self,
        py: Python<'_>,
        args: &PyTuple,
        kwargs: Option<&PyDict>,
    ) -> PyResult<Py<PyAny>> {
        py.allow_threads(|| {
            let mut count = self.count.lock().unwrap();
            *count += 1;
        });
        return self.wraps.call(py, args, kwargs);
    }

2. I would suggest using AtomicUsize instead of Arc<Mutex<usize>> here to avoid locking altogether :)

[zachcoleman]

For 1. thanks! I did not try this combination of allow_threads and Arc+Mutex. This is super cool and clean! Also, a great explanation!

[mejrs]

i.e. call is not atomic and so two calls to this method are happening at the same time

I think that this is the case, and that the executor is yielding in the middle of the hello function. If another function then goes through __call__ and attempts to mutably borrow the decorator, it detects the outstanding borrow. This is not a problem with the changes you have made, because having multiple &self borrows is fine.

[mejrs]

I agree with the suggestion to use an atomic integer here. I'll probably update the example to do exactly that.

[zachcoleman]

Thanks for confirmation and additional explanation!

[zachcoleman]

@mejrs and @davidhewitt thanks for all the help.

I do have a follow-up question... when trying to switch to the AtomicUsize its unclear to me how this could be implemented in this case with &self and my attempt w/ &mut self again generates errors.

Here's the Rust:

use pyo3::prelude::*;
use pyo3::types::{PyDict, PyTuple};
use std::sync::atomic::AtomicUsize;

#[pyclass(name = "Counter")]
pub struct Counter{
    count: AtomicUsize,
    wraps: Py<PyAny>,
}

#[pymethods]
impl Counter{
    #[new]
    fn __new__(wraps: Py<PyAny>) -> Self {
        Counter { 
            count: AtomicUsize::new(0),
            wraps: wraps,
        }
    }

    fn count(&mut self) -> usize{
        // Also how the heck can I make this a read-only function
        // .into_inner() consumes the atomic
        return self.count.get_mut().clone()
    }

    #[args(args = "*", kwargs = "**")]
    fn __call__(
        &mut self,
        py: Python<'_>,
        args: &PyTuple,
        kwargs: Option<&PyDict>,
    ) -> PyResult<Py<PyAny>> {
        // how in here can I mutate an AtomicUsize
        // w/o the interior mutability and allowing for 
        // &self to avoid `Already Borrowed`
        py.allow_threads(|| {
            *self.count.get_mut() += 1;
        });
        return self.wraps.call(py, args, kwargs);
    }
}

#[pymodule]
pub fn rust_counter(_py: Python<'_>, module: &PyModule) -> PyResult<()> {
    module.add_class::<Counter>()?;
    Ok(())
}

[mejrs]

For this purpose, you can see atomic integers as integers that can be mutated from a &self reference (rather than &mut), just like a Mutex<usize> but without the locking. In this case you want the fetch_add and load methods, not get_mut. The ordering used does not matter here.

[zachcoleman]

Duh! Thank you so much!

[mejrs]

You won't need the allow_threads either by the way.