How can I avoid losing data when accessing a Rust class via ffi?

[LilyAcorn]

I'm working on implementing a python profiler in Rust, using PyEval_SetProfile. As part of the C api here, I'm defining a Rust function that is used as a callback and is passed a python object. I've implemented this python object in Rust and I'm now trying to update one of its attributes from Rust in that callback, but the next time the callback is called, that updated data is no longer there. My best guess is that there's a copy somewhere that means I'm not mutating the same Rust object each time.

#[cfg(not(PyPy))]
#[path = ""]
mod _kolo {
    use pyo3::ffi;
    use pyo3::prelude::*;
    use serde_json::json;

    extern "C" fn profile(
        _obj: *mut ffi::PyObject,
        _frame: *mut ffi::PyFrameObject,
        what: i32,
        _arg: *mut ffi::PyObject,
    ) -> i32 {
        let event = {
            if what == ffi::PyTrace_CALL {
                "call"
            } else if what == ffi::PyTrace_RETURN {
                "return"
            } else {
                return 0;
            }
        };
        let _frame = _frame as *mut ffi::PyObject;
        Python::with_gil(|py| unsafe {
            let obj = PyObject::from_borrowed_ptr_or_err(py, _obj);
            let obj = match obj {
                Ok(obj) => obj,
                Err(err) => {
                    err.restore(py);
                    return -1;
                }
            };
            let obj = obj.extract::<KoloProfiler>(py);
            let mut obj = match obj {
                Ok(obj) => obj,
                Err(err) => {
                    err.restore(py);
                    return -1;
                }
            };

            match obj.process_frame(event) {
                Ok(_) => 0,
                Err(err) => {
                    err.restore(py);
                    -1
                }
            }
        })
    }

    #[pyclass(module = "kolo._kolo")]
    #[derive(Clone)]
    struct KoloProfiler {
        frames: Vec<serde_json::Value>,
    }

    impl KoloProfiler {
        fn process_frame(&mut self, event: &str) -> Result<(), PyErr> {
            let frame_data = json!({
                "event": event,
            });
            println!("{}", self.frames.len());  // Always 0
            self.frames.push(frame_data);
            println!("{}", self.frames.len());  // Always 1
            Ok(())
        }
    }

    #[pyfunction]
    fn register_profiler() -> Result<(), PyErr> {
        Python::with_gil(|py| {
            let profiler = KoloProfiler { frames: Vec::new() };
            unsafe {
                ffi::PyEval_SetProfile(Some(profile), profiler.into_py(py).into_ptr());
            }
            Ok(())
        })
    }

    #[pymodule]
    fn _kolo(_py: Python, m: &PyModule) -> PyResult<()> {
        m.add_function(wrap_pyfunction!(register_profiler, m)?)?;
        Ok(())
    }
}

#[cfg(not(PyPy))]
pub use _kolo::*;

When I run this, the length of frames is always 0 or 1 (depending on whether we pushed an item or not), instead of increasing by 1 each time profile calls process_frame. Am I somehow copying the KoloProfiler instance each time profile is called? How can I avoid this?

[mejrs]

When you do let obj = obj.extract::<KoloProfiler>(py);, what is the type of obj here? This (I think) is a KoloProfiler (not &KoloProfiler or &mut KoloProfiler), which means it's cloning it. Try removing that Clone derive - it should raise a compile error there.

What you probably want is to use downcast to get a &PyCell<KoloProfiler> and then borrow_mut to get a &mut KoloProfiler.

(also, using serde_json::Value is probably a bad idea; try using a hashmap or struct instead)

[LilyAcorn]

Thanks! Removing Clone did show a compiler error here.

(I'm using serde_json::Value for a specific reason in the bigger program I extracted this example from. I think it's appropriate there.)

[adamreichold]

#[pyclass] is currently always wrapped inside a PyCell to support interior mutability in the face pervasive shared ownership in Python. The #[derive(Clone)] should not be necessary at all AFAIU. I think you just need to extract PyRefMut<KoloProfiler> instead of a plain KoloProfiler which is a clone of the cell contents, e.g.

let obj = obj.extract::<PyRefMut<KoloProfiler>>(py);

which will call PyCell::try_borrow_mut behind the scenes.

[LilyAcorn]

Yep, PyRefMut was exactly what I needed.