diff --git a/crates/vm/src/stdlib/_thread.rs b/crates/vm/src/stdlib/_thread.rs index 7ac20af92..efe4ed302 100644 --- a/crates/vm/src/stdlib/_thread.rs +++ b/crates/vm/src/stdlib/_thread.rs @@ -104,28 +104,29 @@ pub(crate) mod _thread { } } + fn validate_acquire_args(args: AcquireArgs, vm: &VirtualMachine) -> PyResult<(bool, f64)> { + let timeout = args.timeout.to_secs_f64(); + if !args.blocking && timeout != -1.0 { + return Err(vm.new_value_error("can't specify a timeout for a non-blocking call")); + } + if args.blocking && timeout < 0.0 && timeout != -1.0 { + return Err(vm.new_value_error("timeout value must be a non-negative number")); + } + if timeout > TIMEOUT_MAX { + return Err(vm.new_overflow_error("timeout value is too large")); + } + Ok((args.blocking, timeout)) + } + macro_rules! acquire_lock_impl { - ($mu:expr, $args:expr, $vm:expr) => {{ - let (mu, args, vm) = ($mu, $args, $vm); - let timeout = args.timeout.to_secs_f64(); - match args.blocking { + ($mu:expr, $blocking:expr, $timeout:expr, $vm:expr) => {{ + let (mu, blocking, timeout, vm) = ($mu, $blocking, $timeout, $vm); + match blocking { true if timeout == -1.0 => { vm.allow_threads(|| mu.lock()); Ok(true) } - true if timeout < 0.0 => { - Err(vm.new_value_error("timeout value must be a non-negative number")) - } - true => { - if timeout > TIMEOUT_MAX { - return Err(vm.new_overflow_error("timeout value is too large")); - } - - Ok(vm.allow_threads(|| mu.try_lock_for(Duration::from_secs_f64(timeout)))) - } - false if timeout != -1.0 => { - Err(vm.new_value_error("can't specify a timeout for a non-blocking call")) - } + true => Ok(vm.allow_threads(|| mu.try_lock_for(Duration::from_secs_f64(timeout)))), false => Ok(mu.try_lock()), } }}; @@ -166,7 +167,8 @@ pub(crate) mod _thread { #[pymethod] #[pymethod(name = "acquire_lock")] fn acquire(&self, args: AcquireArgs, vm: &VirtualMachine) -> PyResult { - acquire_lock_impl!(&self.mu, args, vm) + let (blocking, timeout) = validate_acquire_args(args, vm)?; + acquire_lock_impl!(&self.mu, blocking, timeout, vm) } #[pymethod] @@ -252,6 +254,7 @@ pub(crate) mod _thread { #[pymethod] #[pymethod(name = "acquire_lock")] fn acquire(&self, args: AcquireArgs, vm: &VirtualMachine) -> PyResult { + let (blocking, timeout) = validate_acquire_args(args, vm)?; if self.mu.is_owned_by_current_thread() { // Re-entrant acquisition: just increment our count. // parking_lot stays at 1 level; we track recursion ourselves. @@ -259,7 +262,7 @@ pub(crate) mod _thread { .fetch_add(1, core::sync::atomic::Ordering::Relaxed); return Ok(true); } - let result = acquire_lock_impl!(&self.mu, args, vm)?; + let result = acquire_lock_impl!(&self.mu, blocking, timeout, vm)?; if result { self.count.store(1, core::sync::atomic::Ordering::Relaxed); } diff --git a/extra_tests/snippets/builtin_thread.py b/extra_tests/snippets/builtin_thread.py index e9696cce6..3d12e76cc 100644 --- a/extra_tests/snippets/builtin_thread.py +++ b/extra_tests/snippets/builtin_thread.py @@ -1,3 +1,21 @@ import _thread assert _thread.TIMEOUT_MAX in [9223372036.0, 4294967.0] + +from testutils import assert_raises + +# Reentrant calls must validate arguments before incrementing the count. +for error, kwargs in ( + (ValueError, {"blocking": False, "timeout": 1}), + (ValueError, {"timeout": -2}), + (OverflowError, {"timeout": _thread.TIMEOUT_MAX + 1}), +): + lock = _thread.RLock() + assert lock.acquire() + try: + with assert_raises(error): + lock.acquire(**kwargs) + assert lock._recursion_count() == 1 + finally: + while lock._is_owned(): + lock.release()