Skip to content

Commit

Permalink
Add PyPy specific code for CallbackScheduler
Browse files Browse the repository at this point in the history
  • Loading branch information
gi0baro committed Dec 12, 2024
1 parent 0e1c043 commit ad282a5
Show file tree
Hide file tree
Showing 4 changed files with 108 additions and 0 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,9 @@ mimalloc = { version = "0.1.43", default-features = false, features = ["local_dy
[target.'cfg(not(any(target_os = "freebsd", target_os = "windows")))'.dependencies]
tikv-jemallocator = { version = "0.6.0", default-features = false, features = ["disable_initial_exec_tls"] }

[build-dependencies]
pyo3-build-config = { version = "=0.23", features = ["resolve-config"] }

[profile.release]
codegen-units = 1
debug = false
Expand Down
3 changes: 3 additions & 0 deletions build.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
fn main() {
pyo3_build_config::use_pyo3_cfgs();
}
101 changes: 101 additions & 0 deletions src/callbacks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ pub(crate) struct CallbackScheduler {
pyfalse: PyObject,
}

#[cfg(not(PyPy))]
impl CallbackScheduler {
#[inline]
pub(crate) fn schedule(&self, _py: Python, watcher: &PyObject) {
Expand Down Expand Up @@ -113,6 +114,106 @@ impl CallbackScheduler {
}
}

#[cfg(PyPy)]
impl CallbackScheduler {
#[inline]
pub(crate) fn schedule(&self, py: Python, watcher: &PyObject) {
let cbarg = (watcher,).into_pyobject(py).unwrap().into_ptr();
let sched = self.schedule_fn.get().unwrap().as_ptr();

unsafe {
pyo3::ffi::PyObject_CallObject(sched, cbarg);
}
}

#[inline]
pub(crate) fn send(pyself: Py<Self>, py: Python, coro: PyObject) {
let rself = pyself.get();
let ptr = (pyself.clone_ref(py),).into_pyobject(py).unwrap().into_ptr();

unsafe {
pyo3::ffi::PyObject_CallObject(rself.aio_tenter.as_ptr(), ptr);
}

if let Ok(res) = unsafe {
let res = pyo3::ffi::PyObject_CallMethodObjArgs(
coro.as_ptr(),
rself.pyname_aiosend.as_ptr(),
rself.pynone.as_ptr(),
std::ptr::null_mut::<PyObject>(),
);
Bound::from_owned_ptr_or_err(py, res)
} {
if unsafe {
let vptr = pyo3::ffi::PyObject_GetAttr(res.as_ptr(), rself.pyname_aioblock.as_ptr());
Bound::from_owned_ptr_or_err(py, vptr)
.map(|v| v.extract::<bool>().unwrap_or(false))
.unwrap_or(false)
} {
let waker = Py::new(
py,
CallbackSchedulerWaker {
sched: pyself.clone_ref(py),
coro,
},
)
.unwrap();
let resp = res.as_ptr();

unsafe {
pyo3::ffi::PyObject_SetAttr(resp, rself.pyname_aioblock.as_ptr(), rself.pyfalse.as_ptr());
pyo3::ffi::PyObject_CallMethodObjArgs(
resp,
rself.pyname_donecb.as_ptr(),
waker.as_ptr(),
std::ptr::null_mut::<PyObject>(),
);
}
} else {
let sref = Py::new(
py,
CallbackSchedulerRef {
sched: pyself.clone_ref(py),
coro,
},
)
.unwrap();

unsafe {
pyo3::ffi::PyObject_CallMethodObjArgs(
#[allow(clippy::used_underscore_binding)]
rself._loop.as_ptr(),
rself.pyname_loopcs.as_ptr(),
sref.as_ptr(),
std::ptr::null_mut::<PyObject>(),
);
}
}
}

unsafe {
pyo3::ffi::PyObject_CallObject(rself.aio_texit.as_ptr(), ptr);
}
}

#[inline]
pub(crate) fn throw(pyself: Py<Self>, py: Python, coro: PyObject, err: PyObject) {
let pyname_aiothrow = pyself.get().pyname_aiothrow.as_ptr();
let aio_tenter = pyself.get().aio_tenter.as_ptr();
let aio_texit = pyself.get().aio_texit.as_ptr();
let ptr = (pyself,).into_py_any(py).unwrap().as_ptr();
let errptr = (err,).into_py_any(py).unwrap().as_ptr();

unsafe {
let corom = pyo3::ffi::PyObject_GetAttr(coro.as_ptr(), pyname_aiothrow);
pyo3::ffi::PyObject_CallObject(aio_tenter, ptr);
pyo3::ffi::PyObject_CallObject(corom, errptr);
pyo3::ffi::PyErr_Clear();
pyo3::ffi::PyObject_CallObject(aio_texit, ptr);
}
}
}

#[pymethods]
impl CallbackScheduler {
#[new]
Expand Down

0 comments on commit ad282a5

Please sign in to comment.