Merge #2842

2842: Stop leaking in `new_closure` r=adamreichold a=davidhewitt

This is a rebase of #2690 which simplifies the `MaybeLeaked` abstraction from that PR with just `Cow<'static, CStr>`.

This enabled me to annotate with `FIXME` all the places where we still leak; I wonder if we could potentially use `GILOnceCell` in future and statics to avoid those. As those callsities are in `#[pyclass]` and `#[pyfunction]` these are effectively in statics anyway, but it would be nice to tidy up.


Co-authored-by: David Hewitt <1939362+davidhewitt@users.noreply.github.com>

newsfragments/2842.fixed.md

@@ -0,0 +1 @@
+Fix memory leak in `PyCFunction::new_closure`.

src/err/mod.rs

@@ -594,11 +594,6 @@ impl PyErr {
         }
     }
 
-    pub(crate) fn write_unraisable(self, py: Python<'_>, context: PyObject) {
-        self.restore(py);
-        unsafe { ffi::PyErr_WriteUnraisable(context.as_ptr()) };
-    }
-
     #[inline]
     fn from_state(state: PyErrState) -> PyErr {
         PyErr {

src/impl_/pymethods.rs

@@ -1,6 +1,7 @@
-use crate::internal_tricks::{extract_cstr_or_leak_cstring, NulByteInString};
+use crate::exceptions::PyValueError;
 use crate::{ffi, IntoPy, Py, PyAny, PyErr, PyObject, PyResult, PyTraverseError, Python};
-use std::ffi::CStr;
+use std::borrow::Cow;
+use std::ffi::{CStr, CString};
 use std::fmt;
 use std::os::raw::c_int;
 
@@ -95,6 +96,12 @@ pub struct PyClassAttributeDef {
     pub(crate) meth: PyClassAttributeFactory,
 }
 
+impl PyClassAttributeDef {
+    pub(crate) fn attribute_c_string(&self) -> PyResult<Cow<'static, CStr>> {
+        extract_c_string(self.name, "class attribute name cannot contain nul bytes")
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct PyGetterDef {
     pub(crate) name: &'static str,
@@ -161,7 +168,7 @@ impl PyMethodDef {
     }
 
     /// Convert `PyMethodDef` to Python method definition struct `ffi::PyMethodDef`
-    pub(crate) fn as_method_def(&self) -> Result<ffi::PyMethodDef, NulByteInString> {
+    pub(crate) fn as_method_def(&self) -> PyResult<(ffi::PyMethodDef, PyMethodDefDestructor)> {
         let meth = match self.ml_meth {
             PyMethodType::PyCFunction(meth) => ffi::PyMethodDefPointer {
                 PyCFunction: meth.0,
@@ -175,12 +182,16 @@ impl PyMethodDef {
             },
         };
 
-        Ok(ffi::PyMethodDef {
+        let name = get_name(self.ml_name)?;
-            ml_name: get_name(self.ml_name)?.as_ptr(),
+        let doc = get_doc(self.ml_doc)?;
+        let def = ffi::PyMethodDef {
+            ml_name: name.as_ptr(),
             ml_meth: meth,
             ml_flags: self.ml_flags,
-            ml_doc: get_doc(self.ml_doc)?.as_ptr(),
+            ml_doc: doc.as_ptr(),
-        })
+        };
+        let destructor = PyMethodDefDestructor { name, doc };
+        Ok((def, destructor))
     }
 }
 
@@ -214,10 +225,16 @@ impl PyGetterDef {
     /// Copy descriptor information to `ffi::PyGetSetDef`
     pub fn copy_to(&self, dst: &mut ffi::PyGetSetDef) {
         if dst.name.is_null() {
-            dst.name = get_name(self.name).unwrap().as_ptr() as _;
+            let name = get_name(self.name).unwrap();
+            dst.name = name.as_ptr() as _;
+            // FIXME: stop leaking name
+            std::mem::forget(name);
         }
         if dst.doc.is_null() {
-            dst.doc = get_doc(self.doc).unwrap().as_ptr() as _;
+            let doc = get_doc(self.doc).unwrap();
+            dst.doc = doc.as_ptr() as _;
+            // FIXME: stop leaking doc
+            std::mem::forget(doc);
         }
         dst.get = Some(self.meth.0);
     }
@@ -236,21 +253,27 @@ impl PySetterDef {
     /// Copy descriptor information to `ffi::PyGetSetDef`
     pub fn copy_to(&self, dst: &mut ffi::PyGetSetDef) {
         if dst.name.is_null() {
-            dst.name = get_name(self.name).unwrap().as_ptr() as _;
+            let name = get_name(self.name).unwrap();
+            dst.name = name.as_ptr() as _;
+            // FIXME: stop leaking name
+            std::mem::forget(name);
         }
         if dst.doc.is_null() {
-            dst.doc = get_doc(self.doc).unwrap().as_ptr() as _;
+            let doc = get_doc(self.doc).unwrap();
+            dst.doc = doc.as_ptr() as _;
+            // FIXME: stop leaking doc
+            std::mem::forget(doc);
         }
         dst.set = Some(self.meth.0);
     }
 }
 
-fn get_name(name: &'static str) -> Result<&'static CStr, NulByteInString> {
+fn get_name(name: &'static str) -> PyResult<Cow<'static, CStr>> {
-    extract_cstr_or_leak_cstring(name, "Function name cannot contain NUL byte.")
+    extract_c_string(name, "Function name cannot contain NUL byte.")
 }
 
-fn get_doc(doc: &'static str) -> Result<&'static CStr, NulByteInString> {
+fn get_doc(doc: &'static str) -> PyResult<Cow<'static, CStr>> {
-    extract_cstr_or_leak_cstring(doc, "Document cannot contain NUL byte.")
+    extract_c_string(doc, "Document cannot contain NUL byte.")
 }
 
 /// Unwraps the result of __traverse__ for tp_traverse
@@ -263,6 +286,14 @@ pub fn unwrap_traverse_result(result: Result<(), PyTraverseError>) -> c_int {
     }
 }
 
+pub(crate) struct PyMethodDefDestructor {
+    // These members are just to avoid leaking CStrings when possible
+    #[allow(dead_code)]
+    name: Cow<'static, CStr>,
+    #[allow(dead_code)]
+    doc: Cow<'static, CStr>,
+}
+
 // The macros need to Ok-wrap the output of user defined functions; i.e. if they're not a result, make them into one.
 pub trait OkWrap<T> {
     type Error;
@@ -288,3 +319,25 @@ where
         self.map(|o| o.into_py(py))
     }
 }
+
+fn extract_c_string(src: &'static str, err_msg: &'static str) -> PyResult<Cow<'static, CStr>> {
+    let bytes = src.as_bytes();
+    let cow = match bytes {
+        [] => {
+            // Empty string, we can trivially refer to a static "\0" string
+            Cow::Borrowed(unsafe { CStr::from_bytes_with_nul_unchecked(b"\0") })
+        }
+        [.., 0] => {
+            // Last byte is a nul; try to create as a CStr
+            let c_str =
+                CStr::from_bytes_with_nul(bytes).map_err(|_| PyValueError::new_err(err_msg))?;
+            Cow::Borrowed(c_str)
+        }
+        _ => {
+            // Allocate a new CString for this
+            let c_string = CString::new(bytes).map_err(|_| PyValueError::new_err(err_msg))?;
+            Cow::Owned(c_string)
+        }
+    };
+    Ok(cow)
+}

src/internal_tricks.rs

@@ -1,6 +1,4 @@
 use crate::ffi::{Py_ssize_t, PY_SSIZE_T_MAX};
-use std::ffi::{CStr, CString};
-
 pub struct PrivateMarker;
 
 macro_rules! private_decl {
@@ -32,20 +30,6 @@ macro_rules! pyo3_exception {
     };
 }
 
-#[derive(Debug)]
-pub(crate) struct NulByteInString(pub(crate) &'static str);
-
-pub(crate) fn extract_cstr_or_leak_cstring(
-    src: &'static str,
-    err_msg: &'static str,
-) -> Result<&'static CStr, NulByteInString> {
-    CStr::from_bytes_with_nul(src.as_bytes())
-        .or_else(|_| {
-            CString::new(src.as_bytes()).map(|c_string| &*Box::leak(c_string.into_boxed_c_str()))
-        })
-        .map_err(|_| NulByteInString(err_msg))
-}
-
 /// Convert an usize index into a Py_ssize_t index, clamping overflow to
 /// PY_SSIZE_T_MAX.
 pub(crate) fn get_ssize_index(index: usize) -> Py_ssize_t {

src/pyclass.rs

@@ -158,7 +158,12 @@ impl PyTypeBuilder {
             }
             PyMethodDefType::Method(def)
             | PyMethodDefType::Class(def)
-            | PyMethodDefType::Static(def) => self.method_defs.push(def.as_method_def().unwrap()),
+            | PyMethodDefType::Static(def) => {
+                let (def, destructor) = def.as_method_def().unwrap();
+                // FIXME: stop leaking destructor
+                std::mem::forget(destructor);
+                self.method_defs.push(def);
+            }
             // These class attributes are added after the type gets created by LazyStaticType
             PyMethodDefType::ClassAttribute(_) => {}
         }

src/type_object.rs

@@ -2,7 +2,6 @@
 //! Python type object information
 
 use crate::impl_::pyclass::PyClassItemsIter;
-use crate::internal_tricks::extract_cstr_or_leak_cstring;
 use crate::once_cell::GILOnceCell;
 use crate::pyclass::create_type_object;
 use crate::pyclass::PyClass;
@@ -10,6 +9,8 @@ use crate::types::{PyAny, PyType};
 use crate::{conversion::IntoPyPointer, PyMethodDefType};
 use crate::{ffi, AsPyPointer, PyNativeType, PyObject, PyResult, Python};
 use parking_lot::{const_mutex, Mutex};
+use std::borrow::Cow;
+use std::ffi::CStr;
 use std::thread::{self, ThreadId};
 
 /// `T: PyLayout<U>` represents that `T` is a concrete representation of `U` in the Python heap.
@@ -180,11 +181,7 @@ impl LazyStaticType {
         for class_items in items_iter {
             for def in class_items.methods {
                 if let PyMethodDefType::ClassAttribute(attr) = def {
-                    let key = extract_cstr_or_leak_cstring(
+                    let key = attr.attribute_c_string().unwrap();
-                        attr.name,
-                        "class attribute name cannot contain nul bytes",
-                    )
-                    .unwrap();
 
                     match (attr.meth.0)(py) {
                         Ok(val) => items.push((key, val)),
@@ -221,7 +218,7 @@ impl LazyStaticType {
 fn initialize_tp_dict(
     py: Python<'_>,
     type_object: *mut ffi::PyObject,
-    items: Vec<(&'static std::ffi::CStr, PyObject)>,
+    items: Vec<(Cow<'static, CStr>, PyObject)>,
 ) -> PyResult<()> {
     // We hold the GIL: the dictionary update can be considered atomic from
     // the POV of other threads.

src/types/function.rs

@@ -1,14 +1,14 @@
 use crate::derive_utils::PyFunctionArguments;
-use crate::exceptions::PyValueError;
+use crate::methods::PyMethodDefDestructor;
-use crate::impl_::panic::PanicTrap;
+use crate::prelude::*;
-use crate::panic::PanicException;
 use crate::{
     ffi,
     impl_::pymethods::{self, PyMethodDef},
-    types, AsPyPointer,
+    types::{PyCapsule, PyDict, PyTuple},
+    AsPyPointer,
 };
-use crate::{prelude::*, GILPool};
+use std::cell::UnsafeCell;
-use std::os::raw::c_void;
+use std::ffi::CStr;
 
 /// Represents a builtin Python function object.
 #[repr(transparent)]
@@ -16,54 +16,6 @@ pub struct PyCFunction(PyAny);
 
 pyobject_native_type_core!(PyCFunction, ffi::PyCFunction_Type, #checkfunction=ffi::PyCFunction_Check);
 
-const CLOSURE_CAPSULE_NAME: &[u8] = b"pyo3-closure\0";
-
-unsafe extern "C" fn run_closure<F, R>(
-    capsule_ptr: *mut ffi::PyObject,
-    args: *mut ffi::PyObject,
-    kwargs: *mut ffi::PyObject,
-) -> *mut ffi::PyObject
-where
-    F: Fn(&types::PyTuple, Option<&types::PyDict>) -> R + Send + 'static,
-    R: crate::callback::IntoPyCallbackOutput<*mut ffi::PyObject>,
-{
-    crate::impl_::trampoline::cfunction_with_keywords(
-        capsule_ptr,
-        args,
-        kwargs,
-        |py, capsule_ptr, args, kwargs| {
-            let boxed_fn: &F = &*(ffi::PyCapsule_GetPointer(
-                capsule_ptr,
-                CLOSURE_CAPSULE_NAME.as_ptr() as *const _,
-            ) as *mut F);
-            let args = py.from_borrowed_ptr::<types::PyTuple>(args);
-            let kwargs = py.from_borrowed_ptr_or_opt::<types::PyDict>(kwargs);
-            crate::callback::convert(py, boxed_fn(args, kwargs))
-        },
-    )
-}
-
-unsafe extern "C" fn drop_closure<F, R>(capsule_ptr: *mut ffi::PyObject)
-where
-    F: Fn(&types::PyTuple, Option<&types::PyDict>) -> R + Send + 'static,
-    R: crate::callback::IntoPyCallbackOutput<*mut ffi::PyObject>,
-{
-    let trap = PanicTrap::new("uncaught panic during drop_closure");
-    let pool = GILPool::new();
-    if let Err(payload) = std::panic::catch_unwind(|| {
-        let boxed_fn: Box<F> = Box::from_raw(ffi::PyCapsule_GetPointer(
-            capsule_ptr,
-            CLOSURE_CAPSULE_NAME.as_ptr() as *const _,
-        ) as *mut F);
-        drop(boxed_fn)
-    }) {
-        let py = pool.python();
-        let err = PanicException::from_panic_payload(payload);
-        err.write_unraisable(py, "when dropping a closure".into_py(py));
-    };
-    trap.disarm();
-}
-
 impl PyCFunction {
     /// Create a new built-in function with keywords (*args and/or **kwargs).
     pub fn new_with_keywords<'a>(
@@ -101,14 +53,14 @@ impl PyCFunction {
     ///
     /// ```
     /// # use pyo3::prelude::*;
-    /// # use pyo3::{py_run, types};
+    /// # use pyo3::{py_run, types::{PyCFunction, PyDict, PyTuple}};
     ///
     /// Python::with_gil(|py| {
-    ///     let add_one = |args: &types::PyTuple, _kwargs: Option<&types::PyDict>| -> PyResult<_> {
+    ///     let add_one = |args: &PyTuple, _kwargs: Option<&PyDict>| -> PyResult<_> {
     ///         let i = args.extract::<(i64,)>()?.0;
     ///         Ok(i+1)
     ///     };
-    ///     let add_one = types::PyCFunction::new_closure(py, None, None, add_one).unwrap();
+    ///     let add_one = PyCFunction::new_closure(py, None, None, add_one).unwrap();
     ///     py_run!(py, add_one, "assert add_one(42) == 43");
     /// });
     /// ```
@@ -116,46 +68,37 @@ impl PyCFunction {
         py: Python<'a>,
         name: Option<&'static str>,
         doc: Option<&'static str>,
-        f: F,
+        closure: F,
     ) -> PyResult<&'a PyCFunction>
     where
-        F: Fn(&types::PyTuple, Option<&types::PyDict>) -> R + Send + 'static,
+        F: Fn(&PyTuple, Option<&PyDict>) -> R + Send + 'static,
         R: crate::callback::IntoPyCallbackOutput<*mut ffi::PyObject>,
     {
-        let function_ptr = Box::into_raw(Box::new(f));
-        let capsule = unsafe {
-            PyObject::from_owned_ptr_or_err(
-                py,
-                ffi::PyCapsule_New(
-                    function_ptr as *mut c_void,
-                    CLOSURE_CAPSULE_NAME.as_ptr() as *const _,
-                    Some(drop_closure::<F, R>),
-                ),
-            )?
-        };
         let method_def = pymethods::PyMethodDef::cfunction_with_keywords(
             name.unwrap_or("pyo3-closure\0"),
             pymethods::PyCFunctionWithKeywords(run_closure::<F, R>),
             doc.unwrap_or("\0"),
         );
-        Self::internal_new_from_pointers(&method_def, py, capsule.as_ptr(), std::ptr::null_mut())
+        let (def, def_destructor) = method_def.as_method_def()?;
-    }
+
+        let capsule = PyCapsule::new(
+            py,
+            ClosureDestructor::<F> {
+                closure,
+                def: UnsafeCell::new(def),
+                def_destructor,
+            },
+            Some(closure_capsule_name().to_owned()),
+        )?;
+
+        // Safety: just created the capsule with type ClosureDestructor<F> above
+        let data = unsafe { capsule.reference::<ClosureDestructor<F>>() };
 
-    #[doc(hidden)]
-    fn internal_new_from_pointers<'py>(
-        method_def: &PyMethodDef,
-        py: Python<'py>,
-        mod_ptr: *mut ffi::PyObject,
-        module_name: *mut ffi::PyObject,
-    ) -> PyResult<&'py Self> {
-        let def = method_def
-            .as_method_def()
-            .map_err(|err| PyValueError::new_err(err.0))?;
         unsafe {
             py.from_owned_ptr_or_err::<PyCFunction>(ffi::PyCFunction_NewEx(
-                Box::into_raw(Box::new(def)),
+                data.def.get(),
-                mod_ptr,
+                capsule.as_ptr(),
-                module_name,
+                std::ptr::null_mut(),
             ))
         }
     }
@@ -173,10 +116,64 @@ impl PyCFunction {
         } else {
             (std::ptr::null_mut(), std::ptr::null_mut())
         };
-        Self::internal_new_from_pointers(method_def, py, mod_ptr, module_name)
+        let (def, destructor) = method_def.as_method_def()?;
+
+        // FIXME: stop leaking the def and destructor
+        let def = Box::into_raw(Box::new(def));
+        std::mem::forget(destructor);
+
+        unsafe {
+            py.from_owned_ptr_or_err::<PyCFunction>(ffi::PyCFunction_NewEx(
+                def,
+                mod_ptr,
+                module_name,
+            ))
+        }
     }
 }
 
+fn closure_capsule_name() -> &'static CStr {
+    // TODO replace this with const CStr once MSRV new enough
+    CStr::from_bytes_with_nul(b"pyo3-closure\0").unwrap()
+}
+
+unsafe extern "C" fn run_closure<F, R>(
+    capsule_ptr: *mut ffi::PyObject,
+    args: *mut ffi::PyObject,
+    kwargs: *mut ffi::PyObject,
+) -> *mut ffi::PyObject
+where
+    F: Fn(&PyTuple, Option<&PyDict>) -> R + Send + 'static,
+    R: crate::callback::IntoPyCallbackOutput<*mut ffi::PyObject>,
+{
+    crate::impl_::trampoline::cfunction_with_keywords(
+        capsule_ptr,
+        args,
+        kwargs,
+        |py, capsule_ptr, args, kwargs| {
+            let boxed_fn: &ClosureDestructor<F> =
+                &*(ffi::PyCapsule_GetPointer(capsule_ptr, closure_capsule_name().as_ptr())
+                    as *mut ClosureDestructor<F>);
+            let args = py.from_borrowed_ptr::<PyTuple>(args);
+            let kwargs = py.from_borrowed_ptr_or_opt::<PyDict>(kwargs);
+            crate::callback::convert(py, (boxed_fn.closure)(args, kwargs))
+        },
+    )
+}
+
+struct ClosureDestructor<F> {
+    closure: F,
+    // Wrapped in UnsafeCell because Python C-API wants a *mut pointer
+    // to this member.
+    def: UnsafeCell<ffi::PyMethodDef>,
+    // Used to destroy the cstrings in `def`, if necessary.
+    #[allow(dead_code)]
+    def_destructor: PyMethodDefDestructor,
+}
+
+// Safety: F is send and none of the fields are ever mutated
+unsafe impl<F: Send> Send for ClosureDestructor<F> {}
+
 /// Represents a Python function object.
 #[repr(transparent)]
 #[cfg(all(not(Py_LIMITED_API), not(all(PyPy, not(Py_3_8)))))]