diff --git a/rust/hg-cpython/src/ref_sharing.rs b/rust/hg-cpython/src/ref_sharing.rs --- a/rust/hg-cpython/src/ref_sharing.rs +++ b/rust/hg-cpython/src/ref_sharing.rs @@ -23,15 +23,33 @@ //! Macros for use in the `hg-cpython` bridge library. use crate::exceptions::AlreadyBorrowed; -use cpython::{PyClone, PyObject, PyResult, Python}; +use cpython::{exc, PyClone, PyErr, PyObject, PyResult, Python}; use std::cell::{Cell, Ref, RefCell, RefMut}; use std::ops::{Deref, DerefMut}; +use std::sync::atomic::{AtomicUsize, Ordering}; /// Manages the shared state between Python and Rust +/// +/// `PySharedState` is owned by `PySharedRefCell`, and is shared across its +/// derived references. The consistency of these references are guaranteed +/// as follows: +/// +/// - The immutability of `py_class!` object fields. Any mutation of +/// `PySharedRefCell` is allowed only through its `borrow_mut()`. +/// - The `py: Python<'_>` token, which makes sure that any data access is +/// synchronized by the GIL. +/// - The `generation` counter, which increments on `borrow_mut()`. `PyLeaked` +/// reference is valid only if the `current_generation()` equals to the +/// `generation` at the time of `leak_immutable()`. #[derive(Debug, Default)] struct PySharedState { leak_count: Cell, mutably_borrowed: Cell, + // The counter variable could be Cell since any operation on + // PySharedState is synchronized by the GIL, but being "atomic" makes + // PySharedState inherently Sync. The ordering requirement doesn't + // matter thanks to the GIL. + generation: AtomicUsize, } // &PySharedState can be Send because any access to inner cells is @@ -54,6 +72,10 @@ impl PySharedState { match self.leak_count.get() { 0 => { self.mutably_borrowed.replace(true); + // Note that this wraps around to the same value if mutably + // borrowed more than usize::MAX times, which wouldn't happen + // in practice. + self.generation.fetch_add(1, Ordering::Relaxed); Ok(PyRefMut::new(py, pyrefmut, self)) } // TODO @@ -118,6 +140,10 @@ impl PySharedState { self.leak_count.replace(count - 1); } } + + fn current_generation(&self, _py: Python) -> usize { + self.generation.load(Ordering::Relaxed) + } } /// `RefCell` wrapper to be safely used in conjunction with `PySharedState`. @@ -308,14 +334,20 @@ macro_rules! py_shared_ref { } /// Manage immutable references to `PyObject` leaked into Python iterators. +/// +/// This reference will be invalidated once the original value is mutably +/// borrowed. pub struct PyLeaked { inner: PyObject, data: Option, py_shared_state: &'static PySharedState, + /// Generation counter of data `T` captured when PyLeaked is created. + generation: usize, } // DO NOT implement Deref for PyLeaked! Dereferencing PyLeaked -// without taking Python GIL wouldn't be safe. +// without taking Python GIL wouldn't be safe. Also, the underling reference +// is invalid if generation != py_shared_state.generation. impl PyLeaked { /// # Safety @@ -331,14 +363,18 @@ impl PyLeaked { inner: inner.clone_ref(py), data: Some(data), py_shared_state, + generation: py_shared_state.current_generation(py), } } /// Immutably borrows the wrapped value. + /// + /// Borrowing fails if the underlying reference has been invalidated. pub fn try_borrow<'a>( &'a self, py: Python<'a>, ) -> PyResult> { + self.validate_generation(py)?; Ok(PyLeakedRef { _py: py, data: self.data.as_ref().unwrap(), @@ -347,12 +383,15 @@ impl PyLeaked { /// Mutably borrows the wrapped value. /// + /// Borrowing fails if the underlying reference has been invalidated. + /// /// Typically `T` is an iterator. If `T` is an immutable reference, /// `get_mut()` is useless since the inner value can't be mutated. pub fn try_borrow_mut<'a>( &'a mut self, py: Python<'a>, ) -> PyResult> { + self.validate_generation(py)?; Ok(PyLeakedRefMut { _py: py, data: self.data.as_mut().unwrap(), @@ -364,6 +403,13 @@ impl PyLeaked { /// Typically `T` is a static reference to a container, and `U` is an /// iterator of that container. /// + /// # Panics + /// + /// Panics if the underlying reference has been invalidated. + /// + /// This is typically called immediately after the `PyLeaked` is obtained. + /// In which case, the reference must be valid and no panic would occur. + /// /// # Safety /// /// The lifetime of the object passed in to the function `f` is cheated. @@ -375,6 +421,11 @@ impl PyLeaked { py: Python, f: impl FnOnce(T) -> U, ) -> PyLeaked { + // Needs to test the generation value to make sure self.data reference + // is still intact. + self.validate_generation(py) + .expect("map() over invalidated leaked reference"); + // f() could make the self.data outlive. That's why map() is unsafe. // In order to make this function safe, maybe we'll need a way to // temporarily restrict the lifetime of self.data and translate the @@ -384,6 +435,18 @@ impl PyLeaked { inner: self.inner.clone_ref(py), data: Some(new_data), py_shared_state: self.py_shared_state, + generation: self.generation, + } + } + + fn validate_generation(&self, py: Python) -> PyResult<()> { + if self.py_shared_state.current_generation(py) == self.generation { + Ok(()) + } else { + Err(PyErr::new::( + py, + "Cannot access to leaked reference after mutation", + )) } } } @@ -582,6 +645,41 @@ mod test { } #[test] + fn test_leaked_borrow_after_mut() { + let (gil, owner) = prepare_env(); + let py = gil.python(); + let leaked = owner.string_shared(py).leak_immutable().unwrap(); + owner.string(py).py_shared_state.leak_count.replace(0); // XXX cheat + owner.string_shared(py).borrow_mut().unwrap().clear(); + owner.string(py).py_shared_state.leak_count.replace(1); // XXX cheat + assert!(leaked.try_borrow(py).is_err()); + } + + #[test] + fn test_leaked_borrow_mut_after_mut() { + let (gil, owner) = prepare_env(); + let py = gil.python(); + let leaked = owner.string_shared(py).leak_immutable().unwrap(); + let mut leaked_iter = unsafe { leaked.map(py, |s| s.chars()) }; + owner.string(py).py_shared_state.leak_count.replace(0); // XXX cheat + owner.string_shared(py).borrow_mut().unwrap().clear(); + owner.string(py).py_shared_state.leak_count.replace(1); // XXX cheat + assert!(leaked_iter.try_borrow_mut(py).is_err()); + } + + #[test] + #[should_panic(expected = "map() over invalidated leaked reference")] + fn test_leaked_map_after_mut() { + let (gil, owner) = prepare_env(); + let py = gil.python(); + let leaked = owner.string_shared(py).leak_immutable().unwrap(); + owner.string(py).py_shared_state.leak_count.replace(0); // XXX cheat + owner.string_shared(py).borrow_mut().unwrap().clear(); + owner.string(py).py_shared_state.leak_count.replace(1); // XXX cheat + let _leaked_iter = unsafe { leaked.map(py, |s| s.chars()) }; + } + + #[test] fn test_borrow_mut_while_leaked() { let (gil, owner) = prepare_env(); let py = gil.python();