Skip to content

Commit 52af9b4

Browse files
committed
optimize Py<T>::drop for the case when attached
1 parent 26d5a23 commit 52af9b4

File tree

2 files changed

+57
-34
lines changed

2 files changed

+57
-34
lines changed

src/instance.rs

Lines changed: 43 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2009,10 +2009,29 @@ where
20092009
#[cfg(feature = "py-clone")]
20102010
impl<T> Clone for Py<T> {
20112011
#[track_caller]
2012+
#[inline]
20122013
fn clone(&self) -> Self {
2013-
unsafe {
2014-
state::register_incref(self.0);
2014+
#[track_caller]
2015+
#[inline]
2016+
fn try_incref(obj: NonNull<ffi::PyObject>) {
2017+
use crate::internal::state::thread_is_attached;
2018+
2019+
if thread_is_attached() {
2020+
// SAFETY: Py_INCREF is safe to call on a valid Python object if the thread is attached.
2021+
unsafe { ffi::Py_INCREF(obj.as_ptr()) }
2022+
} else {
2023+
incref_failed()
2024+
}
2025+
}
2026+
2027+
#[cold]
2028+
#[track_caller]
2029+
fn incref_failed() -> ! {
2030+
panic!("Cannot clone pointer into Python heap without the thread being attached.");
20152031
}
2032+
2033+
try_incref(self.0);
2034+
20162035
Self(self.0, PhantomData)
20172036
}
20182037
}
@@ -2027,10 +2046,30 @@ impl<T> Clone for Py<T> {
20272046
/// is enabled, it will abort the process.
20282047
impl<T> Drop for Py<T> {
20292048
#[track_caller]
2049+
#[inline]
20302050
fn drop(&mut self) {
2031-
unsafe {
2032-
state::register_decref(self.0);
2051+
// non generic inlineable inner function to reduce code bloat
2052+
#[inline]
2053+
fn inner(obj: NonNull<ffi::PyObject>) {
2054+
use crate::internal::state::thread_is_attached;
2055+
2056+
if thread_is_attached() {
2057+
// SAFETY: Py_DECREF is safe to call on a valid Python object if the thread is attached.
2058+
unsafe { ffi::Py_DECREF(obj.as_ptr()) }
2059+
} else {
2060+
drop_slow(obj)
2061+
}
20332062
}
2063+
2064+
#[cold]
2065+
fn drop_slow(obj: NonNull<ffi::PyObject>) {
2066+
// SAFETY: handing ownership of the reference to `register_decref`.
2067+
unsafe {
2068+
state::register_decref(obj);
2069+
}
2070+
}
2071+
2072+
inner(self.0)
20342073
}
20352074
}
20362075

src/internal/state.rs

Lines changed: 14 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ const ATTACH_FORBIDDEN_DURING_TRAVERSE: isize = -1;
3232
/// 2) PyGILState_Check always returns 1 if the sub-interpreter APIs have ever been called,
3333
/// which could lead to incorrect conclusions that the thread is attached.
3434
#[inline(always)]
35-
fn thread_is_attached() -> bool {
35+
pub(crate) fn thread_is_attached() -> bool {
3636
ATTACH_COUNT.try_with(|c| c.get() > 0).unwrap_or(false)
3737
}
3838

@@ -298,44 +298,28 @@ impl Drop for ForbidAttaching {
298298
}
299299
}
300300

301-
/// Increments the reference count of a Python object if the thread is attached. If
302-
/// the thread is not attached, this function will panic.
303-
///
304-
/// # Safety
305-
/// The object must be an owned Python reference.
306-
#[cfg(feature = "py-clone")]
307-
#[track_caller]
308-
pub unsafe fn register_incref(obj: NonNull<ffi::PyObject>) {
309-
if thread_is_attached() {
310-
unsafe { ffi::Py_INCREF(obj.as_ptr()) }
311-
} else {
312-
panic!("Cannot clone pointer into Python heap without the thread being attached.");
313-
}
314-
}
315-
316301
/// Registers a Python object pointer inside the release pool, to have its reference count decreased
317302
/// the next time the thread is attached in pyo3.
318303
///
319304
/// If the thread is attached, the reference count will be decreased immediately instead of being queued
320305
/// for later.
321306
///
322307
/// # Safety
323-
/// The object must be an owned Python reference.
324-
#[track_caller]
308+
/// - The object must be an owned Python reference.
309+
/// - The reference must not be used after calling this function.
310+
#[inline]
325311
pub unsafe fn register_decref(obj: NonNull<ffi::PyObject>) {
326-
if thread_is_attached() {
327-
unsafe { ffi::Py_DECREF(obj.as_ptr()) }
328-
} else {
329-
#[cfg(not(pyo3_disable_reference_pool))]
312+
#[cfg(not(pyo3_disable_reference_pool))]
313+
{
330314
get_pool().register_decref(obj);
331-
#[cfg(all(
332-
pyo3_disable_reference_pool,
333-
not(pyo3_leak_on_drop_without_reference_pool)
334-
))]
335-
{
336-
let _trap = PanicTrap::new("Aborting the process to avoid panic-from-drop.");
337-
panic!("Cannot drop pointer into Python heap without the thread being attached.");
338-
}
315+
}
316+
#[cfg(all(
317+
pyo3_disable_reference_pool,
318+
not(pyo3_leak_on_drop_without_reference_pool)
319+
))]
320+
{
321+
let _trap = PanicTrap::new("Aborting the process to avoid panic-from-drop.");
322+
panic!("Cannot drop pointer into Python heap without the thread being attached.");
339323
}
340324
}
341325

0 commit comments

Comments
 (0)