add ModuleState, core lifecycle callbacks for #[pymodule]s

newsfragments/5600.added.md

@@ -0,0 +1 @@
+added per-module ModuleState, init + free lifecycle hooks and getters / setters in PyModuleMethods

pyo3-macros-backend/src/module.rs

@@ -531,7 +531,8 @@ fn module_initialization(
                 #pyo3_path::impl_::trampoline::module_exec(module, #module_exec)
             }
 
-            static SLOTS: impl_::PyModuleSlots<4> = impl_::PyModuleSlotsBuilder::new()
+            static SLOTS: impl_::PyModuleSlots<5> = impl_::PyModuleSlotsBuilder::new()
+                .with_mod_exec(impl_::pyo3_module_state_init)
                 .with_mod_exec(__pyo3_module_exec)
                 .with_gil_used(#gil_used)
                 .build();

src/impl_.rs

@@ -28,3 +28,5 @@ pub mod pymodule;
 pub mod trampoline;
 pub mod unindent;
 pub mod wrap;
+
+pub(crate) mod pymodule_state;

src/impl_/pymodule.rs

@@ -4,6 +4,7 @@ use std::{
     cell::UnsafeCell,
     ffi::CStr,
     marker::PhantomData,
+    mem::MaybeUninit,
     os::raw::{c_int, c_void},
 };
 
@@ -24,10 +25,12 @@ use std::sync::atomic::{AtomicI64, Ordering};
 
 #[cfg(not(any(PyPy, GraalPy)))]
 use crate::exceptions::PyImportError;
+use crate::exceptions::PyRuntimeError;
+use crate::impl_::trampoline::trampoline;
 use crate::prelude::PyTypeMethods;
 use crate::{
     ffi,
-    impl_::pyfunction::PyFunctionDef,
+    impl_::{pyfunction::PyFunctionDef, pymodule_state::ModuleState},
     sync::PyOnceLock,
     types::{any::PyAnyMethods, dict::PyDictMethods, PyDict, PyModule, PyModuleMethods},
     Bound, Py, PyAny, PyClass, PyResult, PyTypeInfo, Python,
@@ -77,9 +80,11 @@ impl ModuleDef {
         let ffi_def = UnsafeCell::new(ffi::PyModuleDef {
             m_name: name.as_ptr(),
             m_doc: doc.as_ptr(),
+            m_size: std::mem::size_of::<ModuleState>() as _,
             // TODO: would be slightly nicer to use `[T]::as_mut_ptr()` here,
             // but that requires mut ptr deref on MSRV.
             m_slots: slots.0.get() as _,
+            m_free: Some(pyo3_module_state_free),
             ..INIT
         });
 
@@ -308,6 +313,40 @@ impl PyAddToModule for ModuleDef {
     }
 }
 
+/// Called during multi-phase initialization in order to create an instance of
+/// ModuleState on the memory area specific to modules.
+///
+/// Slot: [`Py_mod_exec`]
+///
+/// [`Py_mod_exec`]: https://docs.python.org/3/c-api/module.html#c.Py_mod_exec
+pub unsafe extern "C" fn pyo3_module_state_init(module: *mut ffi::PyObject) -> c_int {
+    unsafe {
+        trampoline(|_| {
+            let state: *mut MaybeUninit<ModuleState> = ffi::PyModule_GetState(module).cast();
+
+            // CPython builtins just assert this, but cross ffi panics are tricky, so we return an
+            // error instead
+            if state.is_null() {
+                return Err(PyRuntimeError::new_err("PyO3 per-module state was null. This is a bug in the Python interpreter runtime."));
+            }
+
+            (*state).write(ModuleState::new());
+
+            Ok(0)
+        })
+    }
+}
+
+/// Called during deallocation of the module object.
+///
+/// Used for the [`m_free`] field of [`PyModuleDef`].
+///
+/// [`m_free`]: https://docs.python.org/3/c-api/module.html#c.PyModuleDef.m_free
+/// [`PyModuleDef`]: https://docs.python.org/3/c-api/module.html#c.PyModuleDef
+pub unsafe extern "C" fn pyo3_module_state_free(module: *mut c_void) {
+    unsafe { ModuleState::pymodule_free_state(module.cast()) };
+}
+
 #[cfg(test)]
 mod tests {
     use std::{borrow::Cow, ffi::CStr, os::raw::c_int};
@@ -386,6 +425,63 @@ mod tests {
         }
     }
 
+    #[test]
+    fn module_state_init() {
+        use super::{pyo3_module_state_init, ModuleState};
+        use crate::{PyAny, PyErr};
+
+        #[derive(Debug, Clone, Copy, PartialEq, Eq)]
+        struct UserState(u64);
+
+        unsafe extern "C" fn state_test(module: *mut ffi::PyObject) -> c_int {
+            unsafe {
+                trampoline::module_exec(module, |module| {
+                    match ModuleState::pymodule_get_state(module.as_ptr()) {
+                        Some(_) => Ok(()),
+                        None => Err(PyErr::new::<PyAny, _>("failed to initialize ModuleState")),
+                    }
+                })
+            }
+        }
+
+        static SLOTS: PyModuleSlots<5> = PyModuleSlotsBuilder::new()
+            .with_gil_used(false)
+            .with_mod_exec(pyo3_module_state_init)
+            .with_mod_exec(state_test)
+            .build();
+        static MODULE_DEF: ModuleDef = ModuleDef::new(
+            c"test_module_state_init",
+            c"This test is for checking PyO3 ModuleState is initialized correctly",
+            &SLOTS,
+        );
+
+        Python::attach(|py| {
+            let mut module = MODULE_DEF
+                .make_module(py)
+                .expect("module to initialize without error")
+                .into_bound(py);
+            let mystate = UserState(42);
+
+            assert_eq!(
+                None,
+                module.state_ref::<UserState>(),
+                "no state has been added yet"
+            );
+
+            assert_eq!(
+                mystate,
+                *module.state_or_init(|| mystate),
+                "added state successfully"
+            );
+
+            assert_eq!(
+                Some(&mystate),
+                module.state_ref::<UserState>(),
+                "previously added state is referenceable"
+            );
+        })
+    }
+
     #[test]
     #[should_panic]
     fn test_module_slots_builder_overflow() {

src/impl_/pymodule_state.rs

@@ -0,0 +1,192 @@
+use std::ptr::NonNull;
+
+use crate::internal::typemap::{CloneAny, TypeMap};
+use crate::types::PyModule;
+use crate::{ffi, Bound};
+
+/// The internal typemap for [`ModuleState`]
+pub type StateMap = TypeMap<dyn CloneAny + Send>;
+
+/// A marker trait for indicating what type level guarantees (and requirements)
+/// are made for PyO3 `PyModule` state types.
+///
+/// In general, a type *must be*
+///
+/// 1. Fully owned (`'static`)
+/// 2. Cloneable (`Clone`)
+/// 3. Sendable (`Send`)
+///
+/// To qualify as `PyModule` state.
+///
+/// This type is automatically implemented for all types that qualify, so no
+/// further action is required.
+pub trait ModuleStateType: Clone + Send {}
+impl<T: Clone + Send> ModuleStateType for T {}
+
+/// Represents a Python module's state.
+///
+/// More precisely, this `struct` resides on the per-module memory area
+/// allocated during the module's creation.
+#[repr(C)]
+#[derive(Debug)]
+pub struct ModuleState {
+    inner: Option<NonNull<StateCapsule>>,
+}
+
+impl ModuleState {
+    /// Create a new, empty [`ModuleState`]
+    pub fn new() -> Self {
+        let boxed = Box::new(StateCapsule::new());
+
+        Self {
+            inner: NonNull::new(Box::into_raw(boxed)),
+        }
+    }
+
+    pub fn state_map_ref(&self) -> &StateMap {
+        &self.inner_ref().sm
+    }
+
+    pub fn state_map_mut(&mut self) -> &mut StateMap {
+        &mut self.inner_mut().sm
+    }
+
+    fn inner_ref(&self) -> &StateCapsule {
+        self.inner
+            .as_ref()
+            .map(|ptr| unsafe { ptr.as_ref() })
+            .expect("BUG: ModuleState.inner should always be Some, except when dropping")
+    }
+
+    fn inner_mut(&mut self) -> &mut StateCapsule {
+        self.inner
+            .as_mut()
+            .map(|ptr| unsafe { ptr.as_mut() })
+            .expect("BUG: ModuleState.inner should always be Some, except when dropping")
+    }
+
+    /// This is the actual [`Drop::drop`] implementation, split out
+    /// so we can run it on the state ptr returned from [`Self::pymodule_get_state`]
+    ///
+    /// While this function does not take a owned `self`, the calling ModuleState
+    /// should not be accessed again
+    ///
+    /// Calling this function multiple times on a single ModuleState is a noop,
+    /// beyond the first
+    unsafe fn drop_impl(&mut self) {
+        if let Some(ptr) = self.inner.take().map(|state| state.as_ptr()) {
+            // SAFETY: This ptr is allocated via Box::new in Self::new, and is
+            // non null
+            unsafe { drop(Box::from_raw(ptr)) }
+        }
+    }
+}
+
+impl ModuleState {
+    /// Fetch the [`ModuleState`] from a bound PyModule, inheriting it's lifetime
+    ///
+    /// ## Panics
+    ///
+    /// This function can panic if called on a PyModule that has not yet been
+    /// initialized
+    pub(crate) fn from_bound<'a>(this: &'a Bound<'_, PyModule>) -> &'a Self {
+        unsafe {
+            Self::pymodule_get_state(this.as_ptr())
+                .map(|ptr| ptr.as_ref())
+                .expect("pyo3 PyModules should always have per-module state")
+        }
+    }
+
+    /// Fetch the [`ModuleState`] mutably from a bound PyModule, inheriting it's
+    /// lifetime
+    ///
+    /// ## Panics
+    ///
+    /// This function can panic if called on a PyModule that has not yet been
+    /// initialized
+    pub(crate) fn from_bound_mut<'a>(this: &'a mut Bound<'_, PyModule>) -> &'a mut Self {
+        unsafe {
+            Self::pymodule_get_state(this.as_ptr())
+                .map(|mut ptr| ptr.as_mut())
+                .expect("pyo3 PyModules should always have per-module state")
+        }
+    }
+
+    /// Associated low level function for retrieving a pyo3 `pymodule`'s state
+    ///
+    /// If this function returns None, it means the underlying C PyModule does
+    /// not have module state.
+    ///
+    /// This function should only be called on a PyModule that is already
+    /// initialized via PyModule_New (or Py_mod_create)
+    pub(crate) unsafe fn pymodule_get_state(module: *mut ffi::PyObject) -> Option<NonNull<Self>> {
+        unsafe {
+            let state: *mut ModuleState = ffi::PyModule_GetState(module).cast();
+
+            match state.is_null() {
+                true => None,
+                false => Some(NonNull::new_unchecked(state)),
+            }
+        }
+    }
+
+    /// Associated low level function for freeing our `pymodule`'s state
+    /// via a ModuleDef's m_free C callback
+    pub(crate) unsafe fn pymodule_free_state(module: *mut ffi::PyObject) {
+        unsafe {
+            if let Some(state) = Self::pymodule_get_state(module) {
+                // SAFETY: this callback is called when python is freeing the
+                // associated PyModule, so we should never be accessed again
+                (*state.as_ptr()).drop_impl()
+            }
+        }
+    }
+}
+
+impl Drop for ModuleState {
+    fn drop(&mut self) {
+        // SAFETY: we're being dropped, so we'll never be accessed again
+        unsafe { self.drop_impl() };
+    }
+}
+
+impl Default for ModuleState {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Inner layout of [`ModuleState`].
+#[derive(Debug, Clone)]
+struct StateCapsule {
+    sm: StateMap,
+}
+
+impl StateCapsule {
+    fn new() -> Self {
+        Self {
+            sm: StateMap::new(),
+        }
+    }
+}
+
+impl Default for StateCapsule {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn type_assertions() {
+        fn is_send<T: Send>(_t: &T) {}
+        fn is_clone<T: Clone>(_t: &T) {}
+
+        let this = StateCapsule::new();
+        is_send(&this);
+        is_clone(&this);
+    }
+}

src/internal.rs

@@ -2,3 +2,4 @@
 
 pub(crate) mod get_slot;
 pub(crate) mod state;
+pub(crate) mod typemap;