pyo3/internal/
state.rs

1//! Interaction with attachment of the current thread to the Python interpreter.
2
3#[cfg(pyo3_disable_reference_pool)]
4use crate::impl_::panic::PanicTrap;
5use crate::{ffi, Python};
6
7use std::cell::Cell;
8#[cfg(not(pyo3_disable_reference_pool))]
9use std::sync::OnceLock;
10#[cfg_attr(pyo3_disable_reference_pool, allow(unused_imports))]
11use std::{mem, ptr::NonNull, sync};
12
13std::thread_local! {
14    /// This is an internal counter in pyo3 monitoring whether this thread is attached to the interpreter.
15    ///
16    /// It will be incremented whenever an AttachGuard is created, and decremented whenever
17    /// they are dropped.
18    ///
19    /// As a result, if this thread is attached to the interpreter, ATTACH_COUNT is greater than zero.
20    ///
21    /// Additionally, we sometimes need to prevent safe access to the Python interpreter,
22    /// e.g. when implementing `__traverse__`, which is represented by a negative value.
23    static ATTACH_COUNT: Cell<isize> = const { Cell::new(0) };
24}
25
26const ATTACH_FORBIDDEN_DURING_TRAVERSE: isize = -1;
27
28/// Checks whether the thread is attached to the Python interpreter.
29///
30/// Note: This uses pyo3's internal count rather than PyGILState_Check for two reasons:
31///  1) for performance
32///  2) PyGILState_Check always returns 1 if the sub-interpreter APIs have ever been called,
33///     which could lead to incorrect conclusions that the thread is attached.
34#[inline(always)]
35fn thread_is_attached() -> bool {
36    ATTACH_COUNT.try_with(|c| c.get() > 0).unwrap_or(false)
37}
38
39/// RAII type that represents thread attachment to the interpreter.
40pub(crate) enum AttachGuard {
41    /// Indicates the thread was already attached when this AttachGuard was acquired.
42    Assumed,
43    /// Indicates that we attached when this AttachGuard was acquired
44    Ensured { gstate: ffi::PyGILState_STATE },
45}
46
47/// Possible error when calling `try_attach()`
48pub(crate) enum AttachError {
49    /// Forbidden during GC traversal.
50    ForbiddenDuringTraverse,
51    /// The interpreter is not initialized.
52    NotInitialized,
53    #[cfg(Py_3_13)]
54    /// The interpreter is finalizing.
55    Finalizing,
56}
57
58impl AttachGuard {
59    /// PyO3 internal API for attaching to the Python interpreter. The public API is Python::attach.
60    ///
61    /// If the thread was already attached via PyO3, this returns
62    /// `AttachGuard::Assumed`. Otherwise, the thread will attach now and
63    /// `AttachGuard::Ensured` will be returned.
64    pub(crate) fn attach() -> Self {
65        match Self::try_attach() {
66            Ok(guard) => guard,
67            Err(AttachError::ForbiddenDuringTraverse) => {
68                panic!("{}", ForbidAttaching::FORBIDDEN_DURING_TRAVERSE)
69            }
70            Err(AttachError::NotInitialized) => {
71                // try to initialize the interpreter and try again
72                crate::interpreter_lifecycle::ensure_initialized();
73                unsafe { Self::do_attach_unchecked() }
74            }
75            #[cfg(Py_3_13)]
76            Err(AttachError::Finalizing) => {
77                panic!("Cannot attach to the Python interpreter while it is finalizing.");
78            }
79        }
80    }
81
82    /// Variant of the above which will will return gracefully if the interpreter cannot be attached to.
83    pub(crate) fn try_attach() -> Result<Self, AttachError> {
84        match ATTACH_COUNT.try_with(|c| c.get()) {
85            Ok(i) if i > 0 => {
86                // SAFETY: We just checked that the thread is already attached.
87                return Ok(unsafe { Self::assume() });
88            }
89            // Cannot attach during GC traversal.
90            Ok(ATTACH_FORBIDDEN_DURING_TRAVERSE) => {
91                return Err(AttachError::ForbiddenDuringTraverse)
92            }
93            // other cases handled below
94            _ => {}
95        }
96
97        // SAFETY: always safe to call this
98        if unsafe { ffi::Py_IsInitialized() } == 0 {
99            return Err(AttachError::NotInitialized);
100        }
101
102        // Calling `PyGILState_Ensure` while finalizing may crash CPython in unpredictable
103        // ways, we'll make a best effort attempt here to avoid that. (There's a time of
104        // check to time-of-use issue, but it's better than nothing.)
105        //
106        // SAFETY: always safe to call this
107        #[cfg(Py_3_13)]
108        if unsafe { ffi::Py_IsFinalizing() } != 0 {
109            // If the interpreter is not initialized, we cannot attach.
110            return Err(AttachError::Finalizing);
111        }
112
113        // SAFETY: We have done everything reasonable to ensure we're in a safe state to
114        // attach to the Python interpreter.
115        Ok(unsafe { Self::do_attach_unchecked() })
116    }
117
118    /// Acquires the `AttachGuard` without performing any state checking.
119    ///
120    /// This can be called in "unsafe" contexts where the normal interpreter state
121    /// checking performed by `AttachGuard::try_attach` may fail. This includes calling
122    /// as part of multi-phase interpreter initialization.
123    ///
124    /// # Safety
125    ///
126    /// The caller must ensure that the Python interpreter is sufficiently initialized
127    /// for a thread to be able to attach to it.
128    pub(crate) unsafe fn attach_unchecked() -> Self {
129        if thread_is_attached() {
130            return unsafe { Self::assume() };
131        }
132
133        unsafe { Self::do_attach_unchecked() }
134    }
135
136    /// Attach to the interpreter, without a fast-path to check if the thread is already attached.
137    #[cold]
138    unsafe fn do_attach_unchecked() -> Self {
139        // SAFETY: interpreter is sufficiently initialized to attach a thread.
140        let gstate = unsafe { ffi::PyGILState_Ensure() };
141        increment_attach_count();
142        // SAFETY: just attached to the interpreter
143        drop_deferred_references(unsafe { Python::assume_attached() });
144        AttachGuard::Ensured { gstate }
145    }
146
147    /// Acquires the `AttachGuard` while assuming that the thread is already attached
148    /// to the interpreter.
149    pub(crate) unsafe fn assume() -> Self {
150        increment_attach_count();
151        // SAFETY: invariant of calling this function
152        drop_deferred_references(unsafe { Python::assume_attached() });
153        AttachGuard::Assumed
154    }
155
156    /// Gets the Python token associated with this [`AttachGuard`].
157    #[inline]
158    pub(crate) fn python(&self) -> Python<'_> {
159        // SAFETY: this guard guarantees the thread is attached
160        unsafe { Python::assume_attached() }
161    }
162}
163
164/// The Drop implementation for `AttachGuard` will decrement the attach count (and potentially detach).
165impl Drop for AttachGuard {
166    fn drop(&mut self) {
167        match self {
168            AttachGuard::Assumed => {}
169            AttachGuard::Ensured { gstate } => unsafe {
170                // Drop the objects in the pool before attempting to release the thread state
171                ffi::PyGILState_Release(*gstate);
172            },
173        }
174        decrement_attach_count();
175    }
176}
177
178#[cfg(not(pyo3_disable_reference_pool))]
179type PyObjVec = Vec<NonNull<ffi::PyObject>>;
180
181#[cfg(not(pyo3_disable_reference_pool))]
182/// Thread-safe storage for objects which were dec_ref while not attached.
183struct ReferencePool {
184    pending_decrefs: sync::Mutex<PyObjVec>,
185}
186
187#[cfg(not(pyo3_disable_reference_pool))]
188impl ReferencePool {
189    const fn new() -> Self {
190        Self {
191            pending_decrefs: sync::Mutex::new(Vec::new()),
192        }
193    }
194
195    fn register_decref(&self, obj: NonNull<ffi::PyObject>) {
196        self.pending_decrefs.lock().unwrap().push(obj);
197    }
198
199    fn drop_deferred_references(&self, _py: Python<'_>) {
200        let mut pending_decrefs = self.pending_decrefs.lock().unwrap();
201        if pending_decrefs.is_empty() {
202            return;
203        }
204
205        let decrefs = mem::take(&mut *pending_decrefs);
206        drop(pending_decrefs);
207
208        for ptr in decrefs {
209            unsafe { ffi::Py_DECREF(ptr.as_ptr()) };
210        }
211    }
212}
213
214#[cfg(not(pyo3_disable_reference_pool))]
215unsafe impl Send for ReferencePool {}
216
217#[cfg(not(pyo3_disable_reference_pool))]
218unsafe impl Sync for ReferencePool {}
219
220#[cfg(not(pyo3_disable_reference_pool))]
221static POOL: OnceLock<ReferencePool> = OnceLock::new();
222
223#[cfg(not(pyo3_disable_reference_pool))]
224fn get_pool() -> &'static ReferencePool {
225    POOL.get_or_init(ReferencePool::new)
226}
227
228#[cfg_attr(pyo3_disable_reference_pool, inline(always))]
229#[cfg_attr(pyo3_disable_reference_pool, allow(unused_variables))]
230fn drop_deferred_references(py: Python<'_>) {
231    #[cfg(not(pyo3_disable_reference_pool))]
232    if let Some(pool) = POOL.get() {
233        pool.drop_deferred_references(py);
234    }
235}
236
237/// A guard which can be used to temporarily detach from the interpreter and restore on `Drop`.
238pub(crate) struct SuspendAttach {
239    count: isize,
240    tstate: *mut ffi::PyThreadState,
241}
242
243impl SuspendAttach {
244    pub(crate) unsafe fn new() -> Self {
245        let count = ATTACH_COUNT.with(|c| c.replace(0));
246        let tstate = unsafe { ffi::PyEval_SaveThread() };
247
248        Self { count, tstate }
249    }
250}
251
252impl Drop for SuspendAttach {
253    fn drop(&mut self) {
254        ATTACH_COUNT.with(|c| c.set(self.count));
255        unsafe {
256            ffi::PyEval_RestoreThread(self.tstate);
257
258            // Update counts of `Py<T>` that were dropped while not attached.
259            #[cfg(not(pyo3_disable_reference_pool))]
260            if let Some(pool) = POOL.get() {
261                pool.drop_deferred_references(Python::assume_attached());
262            }
263        }
264    }
265}
266
267/// Used to lock safe access to the interpreter
268pub(crate) struct ForbidAttaching {
269    count: isize,
270}
271
272impl ForbidAttaching {
273    const FORBIDDEN_DURING_TRAVERSE: &'static str = "Attaching a thread to the interpreter is prohibited while a __traverse__ implementation is running.";
274
275    /// Lock access to the interpreter while an implementation of `__traverse__` is running
276    pub fn during_traverse() -> Self {
277        Self::new(ATTACH_FORBIDDEN_DURING_TRAVERSE)
278    }
279
280    fn new(reason: isize) -> Self {
281        let count = ATTACH_COUNT.with(|c| c.replace(reason));
282
283        Self { count }
284    }
285
286    #[cold]
287    fn bail(current: isize) {
288        match current {
289            ATTACH_FORBIDDEN_DURING_TRAVERSE => panic!("{}", Self::FORBIDDEN_DURING_TRAVERSE),
290            _ => panic!("Attaching a thread to the interpreter is currently prohibited."),
291        }
292    }
293}
294
295impl Drop for ForbidAttaching {
296    fn drop(&mut self) {
297        ATTACH_COUNT.with(|c| c.set(self.count));
298    }
299}
300
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]
308pub 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
316/// Registers a Python object pointer inside the release pool, to have its reference count decreased
317/// the next time the thread is attached in pyo3.
318///
319/// If the thread is attached, the reference count will be decreased immediately instead of being queued
320/// for later.
321///
322/// # Safety
323/// The object must be an owned Python reference.
324#[track_caller]
325pub 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))]
330        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        }
339    }
340}
341
342/// Private helper function to check if we are currently in a GC traversal (as detected by PyO3).
343#[cfg(any(not(Py_LIMITED_API), Py_3_11))]
344pub(crate) fn is_in_gc_traversal() -> bool {
345    ATTACH_COUNT
346        .try_with(|c| c.get() == ATTACH_FORBIDDEN_DURING_TRAVERSE)
347        .unwrap_or(false)
348}
349
350/// Increments pyo3's internal attach count - to be called whenever an AttachGuard is created.
351#[inline(always)]
352fn increment_attach_count() {
353    // Ignores the error in case this function called from `atexit`.
354    let _ = ATTACH_COUNT.try_with(|c| {
355        let current = c.get();
356        if current < 0 {
357            ForbidAttaching::bail(current);
358        }
359        c.set(current + 1);
360    });
361}
362
363/// Decrements pyo3's internal attach count - to be called whenever AttachGuard is dropped.
364#[inline(always)]
365fn decrement_attach_count() {
366    // Ignores the error in case this function called from `atexit`.
367    let _ = ATTACH_COUNT.try_with(|c| {
368        let current = c.get();
369        debug_assert!(
370            current > 0,
371            "Negative attach count detected. Please report this error to the PyO3 repo as a bug."
372        );
373        c.set(current - 1);
374    });
375}
376
377#[cfg(test)]
378mod tests {
379    use super::*;
380
381    use crate::{ffi, types::PyAnyMethods, Py, PyAny, Python};
382
383    fn get_object(py: Python<'_>) -> Py<PyAny> {
384        py.eval(ffi::c_str!("object()"), None, None)
385            .unwrap()
386            .unbind()
387    }
388
389    #[cfg(not(pyo3_disable_reference_pool))]
390    fn pool_dec_refs_does_not_contain(obj: &Py<PyAny>) -> bool {
391        !get_pool()
392            .pending_decrefs
393            .lock()
394            .unwrap()
395            .contains(&unsafe { NonNull::new_unchecked(obj.as_ptr()) })
396    }
397
398    // With free-threading, threads can empty the POOL at any time, so this
399    // function does not test anything meaningful
400    #[cfg(not(any(pyo3_disable_reference_pool, Py_GIL_DISABLED)))]
401    fn pool_dec_refs_contains(obj: &Py<PyAny>) -> bool {
402        get_pool()
403            .pending_decrefs
404            .lock()
405            .unwrap()
406            .contains(&unsafe { NonNull::new_unchecked(obj.as_ptr()) })
407    }
408
409    #[test]
410    fn test_pyobject_drop_attached_decreases_refcnt() {
411        Python::attach(|py| {
412            let obj = get_object(py);
413
414            // Create a reference to drop while attached.
415            let reference = obj.clone_ref(py);
416
417            assert_eq!(obj.get_refcnt(py), 2);
418            #[cfg(not(pyo3_disable_reference_pool))]
419            assert!(pool_dec_refs_does_not_contain(&obj));
420
421            // While attached, reference count will be decreased immediately.
422            drop(reference);
423
424            assert_eq!(obj.get_refcnt(py), 1);
425            #[cfg(not(any(pyo3_disable_reference_pool)))]
426            assert!(pool_dec_refs_does_not_contain(&obj));
427        });
428    }
429
430    #[test]
431    #[cfg(all(not(pyo3_disable_reference_pool), not(target_arch = "wasm32")))] // We are building wasm Python with pthreads disabled
432    fn test_pyobject_drop_detached_doesnt_decrease_refcnt() {
433        let obj = Python::attach(|py| {
434            let obj = get_object(py);
435            // Create a reference to drop while detached.
436            let reference = obj.clone_ref(py);
437
438            assert_eq!(obj.get_refcnt(py), 2);
439            assert!(pool_dec_refs_does_not_contain(&obj));
440
441            // Drop reference in a separate (detached) thread.
442            std::thread::spawn(move || drop(reference)).join().unwrap();
443
444            // The reference count should not have changed, it is remembered
445            // to release later.
446            assert_eq!(obj.get_refcnt(py), 2);
447            #[cfg(not(Py_GIL_DISABLED))]
448            assert!(pool_dec_refs_contains(&obj));
449            obj
450        });
451
452        // On next attach, the reference is released
453        #[allow(unused)]
454        Python::attach(|py| {
455            // With free-threading, another thread could still be processing
456            // DECREFs after releasing the lock on the POOL, so the
457            // refcnt could still be 2 when this assert happens
458            #[cfg(not(Py_GIL_DISABLED))]
459            assert_eq!(obj.get_refcnt(py), 1);
460            assert!(pool_dec_refs_does_not_contain(&obj));
461        });
462    }
463
464    #[test]
465    #[allow(deprecated)]
466    fn test_attach_counts() {
467        // Check `attach` and AttachGuard both increase counts correctly
468        let get_attach_count = || ATTACH_COUNT.with(|c| c.get());
469
470        assert_eq!(get_attach_count(), 0);
471        Python::attach(|_| {
472            assert_eq!(get_attach_count(), 1);
473
474            let pool = unsafe { AttachGuard::assume() };
475            assert_eq!(get_attach_count(), 2);
476
477            let pool2 = unsafe { AttachGuard::assume() };
478            assert_eq!(get_attach_count(), 3);
479
480            drop(pool);
481            assert_eq!(get_attach_count(), 2);
482
483            Python::attach(|_| {
484                // nested `attach` updates attach count
485                assert_eq!(get_attach_count(), 3);
486            });
487            assert_eq!(get_attach_count(), 2);
488
489            drop(pool2);
490            assert_eq!(get_attach_count(), 1);
491        });
492        assert_eq!(get_attach_count(), 0);
493    }
494
495    #[test]
496    fn test_detach() {
497        assert!(!thread_is_attached());
498
499        Python::attach(|py| {
500            assert!(thread_is_attached());
501
502            py.detach(move || {
503                assert!(!thread_is_attached());
504
505                Python::attach(|_| assert!(thread_is_attached()));
506
507                assert!(!thread_is_attached());
508            });
509
510            assert!(thread_is_attached());
511        });
512
513        assert!(!thread_is_attached());
514    }
515
516    #[cfg(feature = "py-clone")]
517    #[test]
518    #[should_panic]
519    fn test_detach_updates_refcounts() {
520        Python::attach(|py| {
521            // Make a simple object with 1 reference
522            let obj = get_object(py);
523            assert!(obj.get_refcnt(py) == 1);
524            // Cloning the object when detached should panic
525            py.detach(|| obj.clone());
526        });
527    }
528
529    #[test]
530    fn recursive_attach_ok() {
531        Python::attach(|py| {
532            let obj = Python::attach(|_| py.eval(ffi::c_str!("object()"), None, None).unwrap());
533            assert_eq!(obj.get_refcnt(), 1);
534        })
535    }
536
537    #[cfg(feature = "py-clone")]
538    #[test]
539    fn test_clone_attached() {
540        Python::attach(|py| {
541            let obj = get_object(py);
542            let count = obj.get_refcnt(py);
543
544            // Cloning when attached should increase reference count immediately
545            #[allow(clippy::redundant_clone)]
546            let c = obj.clone();
547            assert_eq!(count + 1, c.get_refcnt(py));
548        })
549    }
550
551    #[test]
552    #[cfg(not(pyo3_disable_reference_pool))]
553    fn test_drop_deferred_references_does_not_deadlock() {
554        // drop_deferred_references can run arbitrary Python code during Py_DECREF.
555        // if the locking is implemented incorrectly, it will deadlock.
556
557        use crate::ffi;
558
559        Python::attach(|py| {
560            let obj = get_object(py);
561
562            unsafe extern "C" fn capsule_drop(capsule: *mut ffi::PyObject) {
563                // This line will implicitly call drop_deferred_references
564                // -> and so cause deadlock if drop_deferred_references is not handling recursion correctly.
565                let pool = unsafe { AttachGuard::assume() };
566
567                // Rebuild obj so that it can be dropped
568                unsafe {
569                    Py::<PyAny>::from_owned_ptr(
570                        pool.python(),
571                        ffi::PyCapsule_GetPointer(capsule, std::ptr::null()) as _,
572                    )
573                };
574            }
575
576            let ptr = obj.into_ptr();
577
578            let capsule =
579                unsafe { ffi::PyCapsule_New(ptr as _, std::ptr::null(), Some(capsule_drop)) };
580
581            get_pool().register_decref(NonNull::new(capsule).unwrap());
582
583            // Updating the counts will call decref on the capsule, which calls capsule_drop
584            get_pool().drop_deferred_references(py);
585        })
586    }
587
588    #[test]
589    #[cfg(not(pyo3_disable_reference_pool))]
590    fn test_attach_guard_drop_deferred_references() {
591        Python::attach(|py| {
592            let obj = get_object(py);
593
594            // For AttachGuard::attach
595
596            get_pool().register_decref(NonNull::new(obj.clone_ref(py).into_ptr()).unwrap());
597            #[cfg(not(Py_GIL_DISABLED))]
598            assert!(pool_dec_refs_contains(&obj));
599            let _guard = AttachGuard::attach();
600            assert!(pool_dec_refs_does_not_contain(&obj));
601
602            // For AttachGuard::assume
603
604            get_pool().register_decref(NonNull::new(obj.clone_ref(py).into_ptr()).unwrap());
605            #[cfg(not(Py_GIL_DISABLED))]
606            assert!(pool_dec_refs_contains(&obj));
607            let _guard2 = unsafe { AttachGuard::assume() };
608            assert!(pool_dec_refs_does_not_contain(&obj));
609        })
610    }
611}
pyo3/internal/state.rs

pyo3/internal/
state.rs
