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 <[email protected]>

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 {
-            ml_name: get_name(self.ml_name)?.as_ptr(),
+        let name = get_name(self.ml_name)?;
+        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> {
-    extract_cstr_or_leak_cstring(name, "Function name cannot contain NUL byte.")
+fn get_name(name: &'static str) -> PyResult<Cow<'static, CStr>> {
+    extract_c_string(name, "Function name cannot contain NUL byte.")
 }
 
-fn get_doc(doc: &'static str) -> Result<&'static CStr, NulByteInString> {
-    extract_cstr_or_leak_cstring(doc, "Document cannot contain NUL byte.")
+fn get_doc(doc: &'static str) -> PyResult<Cow<'static, CStr>> {
+    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,14 +2,15 @@
 //! 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;
 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(
-                        attr.name,
-                        "class attribute name cannot contain nul bytes",
-                    )
-                    .unwrap();
+                    let key = attr.attribute_c_string().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.