Skip to content
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
123 changes: 85 additions & 38 deletions crates/stdlib/src/_queue.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,9 @@ mod _queue {
use std::time::Instant;

use crate::vm::{
AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine,
builtins::{PyBaseExceptionRef, PyException, PyGenericAlias, PyStr, PyType, PyTypeRef},
function::TimeoutSeconds,
function::{ArgIntoFloat, OptionalArg},
protocol::PyNumberMethods,
types::{AsNumber, Constructor, Representable},
};
Expand Down Expand Up @@ -85,15 +85,6 @@ mod _queue {
vm.allow_threads(|| self.mutex.lock())
}

fn release(&self, vm: &VirtualMachine) {
{
let mut count = self.lock_count(vm);
*count += 1;
} // lock dropped. now we can notify a waiting thread

self.cond.notify_one();
}

/// `Ok(true)` if acquired, `Ok(false)` on timeout, `Err` if a signal
/// handler raised (e.g. `KeyboardInterrupt`) while we were waiting.
fn acquire(
Expand Down Expand Up @@ -167,10 +158,21 @@ mod _queue {
expect(unused_variables, reason = "only the semaphore needs the vm")
)]
fn push(&self, item: PyObjectRef, vm: &VirtualMachine) {
self.buf.lock().push_back(item);

#[cfg(feature = "threading")]
self.sem.release(vm);
{
// Publish the permit before unlocking the buffer, so another
// thread cannot observe an item that no getter can claim.
let mut count = self.sem.lock_count(vm);
{
let mut buf = self.buf.lock();
buf.push_back(item);
*count += 1;
}
drop(count);
self.sem.cond.notify_one();
}
#[cfg(not(feature = "threading"))]
self.buf.lock().push_back(item);
}

/// Returns a strong reference from the head of the buffer.
Expand Down Expand Up @@ -204,26 +206,25 @@ mod _queue {
struct PutArgs {
#[pyarg(any)]
item: PyObjectRef,
#[expect(
dead_code,
reason = "Intentional. Provide compatibility with the Queue class"
)]
#[pyarg(any, default = true)]
block: bool,
#[expect(
dead_code,
reason = "Intentional. Provide compatibility with the Queue class"
)]
#[pyarg(any, optional, py_default = "True")]
block: OptionalArg<PyObjectRef>,
#[pyarg(any, optional)]
timeout: Option<PyObjectRef>,
}

#[derive(FromArgs)]
struct GetArgs {
#[pyarg(any, default = true)]
block: bool,
#[pyarg(any, optional, py_default = "True")]
block: OptionalArg<PyObjectRef>,
#[pyarg(any, optional)]
timeout: Option<TimeoutSeconds>,
timeout: Option<PyObjectRef>,
}

fn is_blocking(block: &OptionalArg<PyObjectRef>, vm: &VirtualMachine) -> PyResult<bool> {
match block {
OptionalArg::Present(value) => value.try_to_bool(vm),
OptionalArg::Missing => Ok(true),
}
}

#[pyclass(
Expand All @@ -248,9 +249,16 @@ mod _queue {
}

#[pymethod]
fn put(zelf: &Py<Self>, args: PutArgs, vm: &VirtualMachine) {
let PutArgs { item, .. } = args;
fn put(zelf: &Py<Self>, args: PutArgs, vm: &VirtualMachine) -> PyResult<()> {
let PutArgs {
item,
block,
timeout,
} = args;
is_blocking(&block, vm)?;
zelf.push(item, vm);
drop(timeout);
Ok(())
}

#[pymethod]
Expand All @@ -261,10 +269,11 @@ mod _queue {
#[pymethod]
fn get(zelf: &Py<Self>, args: GetArgs, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
let GetArgs { block, timeout } = args;
let blocking = is_blocking(&block, vm)?;

// Non-blocking: just try once
if !block {
return Self::get_inner(&mut zelf.buf.lock()).ok_or_else(|| empty_error(vm));
if !blocking {
return Self::get_nowait(zelf, vm);
}

#[cfg_attr(
Expand All @@ -274,22 +283,56 @@ mod _queue {
reason = "We are still validating the 'timeout' arg even if we don't have threading"
)
)]
let deadline = match timeout.map(|v| v.to_secs_f64()) {
Some(v) if v < 0.0 => {
return Err(vm.new_value_error("'timeout' must be a non-negative number"));
let deadline = match timeout.as_ref() {
Some(value) => {
let overflow =
|| vm.new_overflow_error("timestamp out of range for C PyTime_t");
let nanos = if let Some(index) = value.try_index_opt(vm) {
let index = index.map_err(|error| {
if error.fast_isinstance(vm.ctx.exceptions.overflow_error) {
overflow()
} else {
error
}
})?;
let seconds: i64 = index.try_to_primitive(vm).map_err(|_| overflow())?;
seconds.checked_mul(1_000_000_000).ok_or_else(overflow)?
} else {
let seconds =
ArgIntoFloat::try_from_object(vm, value.clone())?.into_float();
if seconds.is_nan() {
return Err(vm.new_value_error("Invalid value NaN (not a number)"));
}
let nanos = (seconds * 1_000_000_000.0).ceil();
if !(i64::MIN as f64..-(i64::MIN as f64)).contains(&nanos) {
return Err(overflow());
}
nanos as i64
};
if nanos < 0 {
return Err(vm.new_value_error("'timeout' must be a non-negative number"));
}
Some(
Instant::now()
.checked_add(Duration::from_nanos(nanos as u64))
.ok_or_else(overflow)?,
)
}
Some(v) => Some(Instant::now() + Duration::from_secs_f64(v)),
None => None,
};

#[cfg(feature = "threading")]
{
if !zelf.sem.acquire(block, deadline, vm)? {
if !zelf.sem.acquire(blocking, deadline, vm)? {
return Err(empty_error(vm));
}
}

Self::get_inner(&mut zelf.buf.lock()).ok_or_else(|| empty_error(vm))
let item = {
let mut buf = zelf.buf.lock();
Self::get_inner(&mut buf)
};
item.ok_or_else(|| empty_error(vm))
}

#[pymethod]
Expand All @@ -301,7 +344,11 @@ mod _queue {
}
}

Self::get_inner(&mut zelf.buf.lock()).ok_or_else(|| empty_error(vm))
let item = {
let mut buf = zelf.buf.lock();
Self::get_inner(&mut buf)
};
item.ok_or_else(|| empty_error(vm))
}

#[pyclassmethod]
Expand Down
Loading