Skip to content

Share ctypes backing-owner snapshots between aggregate copies #9392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
youknowone merged 3 commits into from
Oct 27, 2026
Merged

Share ctypes backing-owner snapshots between aggregate copies #9392

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
986e6a3
Share ctypes backing-owner snapshots between aggregate copies
1ndahous3 Oct 26, 2026
332101a
Limit ctypes snapshot helper to supported targets
1ndahous3 Oct 26, 2026
f4df593
Clarify ctypes snapshot allocation helper names
1ndahous3 Oct 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 15 additions & 2 deletions crates/vm/src/gc_state.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1373,8 +1373,19 @@ pub fn current_owner() -> GcOwner {
/// it collect if the allocation pushed gen0 past its threshold.
///
/// # Safety
/// obj must be a valid pointer to a PyObject that is not already tracked.
/// obj must be a valid pointer to an untracked PyObject with freshly zeroed `gc_bits`.
#[inline]
pub(crate) unsafe fn track_new_object(obj: NonNull<PyObject>) {
unsafe { track_new_object_impl(obj, true) };
}

/// Track fresh internal storage, optionally allowing collection before its
/// caller publishes it.
///
/// # Safety
/// obj must be a valid pointer to an untracked PyObject with freshly zeroed `gc_bits`.
#[inline]
pub(crate) unsafe fn track_new_object_impl(obj: NonNull<PyObject>, allow_collection: bool) {
let state = gc_state();
let Some(gc) = crate::vm::thread::current_gc_state() else {
// No interpreter is running: the shared context builds its own objects
Expand All @@ -1385,7 +1396,9 @@ pub(crate) unsafe fn track_new_object(obj: NonNull<PyObject>) {
// SAFETY: as in `current_owner`.
let gc = unsafe { gc.as_ref() };
unsafe { state.track_object_fresh(obj, gc.owner) };
state.maybe_collect(gc);
if allow_collection {
state.maybe_collect(gc);
}
}

/// Track a generator (or coroutine, or async generator) together with the
Expand Down
36 changes: 33 additions & 3 deletions crates/vm/src/object/core.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2946,6 +2946,36 @@ impl<T: PyPayload> PyRef<T> {
impl<T: PyPayload + crate::object::MaybeTraverse + core::fmt::Debug> PyRef<T> {
#[inline(always)]
pub fn new_ref(payload: T, typ: crate::builtins::PyTypeRef, dict: Option<PyDictRef>) -> Self {
Self::new_ref_impl(payload, typ, dict, true)
}

/// Allocate and track internal storage without running collection during
/// publication. The next ordinary tracked allocation still checks thresholds.
#[cfg(all(
feature = "host_env",
any(
target_os = "linux",
target_os = "macos",
target_os = "windows",
target_os = "android"
),
not(any(target_env = "musl", target_env = "sgx"))
))]
pub(crate) fn new_ref_without_collection(

@youknowone youknowone Oct 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the naming is confusing. new_ref_without_collection is default false call of new_ref_with_collection. those names are looking symmetric.

payload: T,
typ: crate::builtins::PyTypeRef,
dict: Option<PyDictRef>,
) -> Self {
Self::new_ref_impl(payload, typ, dict, false)
}

#[inline(always)]
fn new_ref_impl(
payload: T,
typ: crate::builtins::PyTypeRef,
dict: Option<PyDictRef>,
allow_collection: bool,
) -> Self {
let has_dict = dict.is_some();
let is_heaptype = typ.heaptype_ext.is_some();
let track = !T::NEW_REF_UNTRACKED
Expand Down Expand Up @@ -2988,10 +3018,10 @@ impl<T: PyPayload + crate::object::MaybeTraverse + core::fmt::Debug> PyRef<T> {
// unless the payload opts out via NEW_REF_UNTRACKED (e.g. call frames,
// which are tracked lazily only on escape).
if track {
// Tracks under the interpreter running now and collects if this
// allocation pushed gen0 past its threshold.
// Tracks under the interpreter running now and, when allowed,
// collects if this allocation pushed gen0 past its threshold.
unsafe {
crate::gc_state::track_new_object(ptr.cast());
crate::gc_state::track_new_object_impl(ptr.cast(), allow_collection);
}
}

Expand Down
186 changes: 141 additions & 45 deletions crates/vm/src/stdlib/_ctypes/base.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -500,48 +500,94 @@ unsafe impl Traverse for PyCData {
}
}

#[derive(Debug, Default)]
struct PrivateRefs {
slots: std::collections::HashMap<String, Vec<usize>, rapidhash::quality::RandomState>,
owners: std::collections::HashMap<usize, (PyObjectRef, usize), rapidhash::quality::RandomState>,
#[derive(Debug, Clone)]
enum HiddenRefs {
Single(PyObjectRef),
Snapshot(crate::builtins::PyTupleRef),
}

// SAFETY: the owner map stores one strong reference per identity; slots hold IDs.
unsafe impl Traverse for PrivateRefs {
#[expect(clippy::iter_over_hash_type, reason = "GC edge order is irrelevant")]
fn traverse(&self, tracer: &mut TraverseFn<'_>) {
for (owner, _) in self.owners.values() {
owner.traverse(tracer);
impl HiddenRefs {
fn id(&self) -> usize {
match self {
Self::Single(obj) => obj.get_id(),
Self::Snapshot(tuple) => tuple.get_id(),
}
}
fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
self.slots.clear();
out.extend(self.owners.drain().map(|(_, (owner, _))| owner));

fn owners(&self) -> &[PyObjectRef] {
match self {
Self::Single(obj) => core::slice::from_ref(obj),
Self::Snapshot(tuple) => tuple.as_slice(),
}
}

fn into_object(self) -> PyObjectRef {
match self {
Self::Single(obj) => obj,
Self::Snapshot(tuple) => tuple.into(),
}
}
}

#[derive(Debug, Default)]
struct PrivateRefs {
slots: std::collections::HashMap<String, usize, rapidhash::quality::RandomState>,
groups: std::collections::HashMap<usize, (HiddenRefs, usize), rapidhash::quality::RandomState>,
owners: std::collections::HashMap<usize, (PyObjectRef, usize), rapidhash::quality::RandomState>,
snapshot: Option<crate::builtins::PyTupleRef>,
}

impl PrivateRefs {
fn replace(&mut self, key: String, owners: Vec<PyObjectRef>) -> Vec<PyObjectRef> {
let old = if owners.is_empty() {
self.slots.remove(&key)
fn replace(&mut self, key: String, incoming: Option<HiddenRefs>) -> Vec<PyObjectRef> {
let old = if let Some(group) = &incoming {
self.slots.insert(key, group.id())
} else {
self.slots
.insert(key, owners.iter().map(|obj| obj.get_id()).collect())
self.slots.remove(&key)
};
// Count each slot's references, but retain each Python owner only once.
// Add first so replacing an owner with itself cannot release it.
for obj in owners {
self.owners.entry(obj.get_id()).or_insert((obj, 0)).1 += 1;
}
let mut displaced = Vec::new();
for id in old.into_iter().flatten() {
let owner = self.owners.get_mut(&id).expect("registered private owner");
owner.1 -= 1;
if owner.1 == 0 {
displaced.push(self.owners.remove(&id).unwrap().0);
let mut changed = false;
if let Some(group) = incoming {
let id = group.id();
if let Some((_, count)) = self.groups.get_mut(&id) {
*count += 1;
displaced.push(group.into_object());
} else {
for obj in group.owners() {
match self.owners.entry(obj.get_id()) {
std::collections::hash_map::Entry::Occupied(mut entry) => {
entry.get_mut().1 += 1
}
std::collections::hash_map::Entry::Vacant(entry) => {
entry.insert((obj.clone(), 1));
changed = true;
}
}
}
self.groups.insert(id, (group, 1));
}
}
// Membership is counted per distinct group, not per destination slot.
// Publish incoming owners before removing old ones, including self-copy.
if let Some(id) = old {
let (_, count) = self.groups.get_mut(&id).expect("registered private group");
*count -= 1;
if *count == 0 {
let (group, _) = self.groups.remove(&id).unwrap();
for obj in group.owners() {
let id = obj.get_id();
let (_, count) = self.owners.get_mut(&id).expect("registered private owner");
*count -= 1;
if *count == 0 {
displaced.push(self.owners.remove(&id).unwrap().0);
changed = true;
}
}
displaced.push(group.into_object());
}
}
// The caller releases these after unlocking and publishing _objects.
if changed {
displaced.extend(self.snapshot.take().map(Into::into));
}
displaced
}

Expand All @@ -550,6 +596,34 @@ impl PrivateRefs {
}
}

// SAFETY: visit each stored strong reference once. Shared snapshot elements
// belong to the Python tuple and are traced by that tuple, not by each slot.
unsafe impl Traverse for PrivateRefs {
#[expect(clippy::iter_over_hash_type, reason = "GC edge order is irrelevant")]
fn traverse(&self, tracer: &mut TraverseFn<'_>) {
for (obj, _) in self.owners.values() {
obj.traverse(tracer);
}
for (group, _) in self.groups.values() {
match group {
HiddenRefs::Single(obj) => obj.traverse(tracer),
HiddenRefs::Snapshot(tuple) => tuple.traverse(tracer),
}
}
self.snapshot.traverse(tracer);
}
fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
self.slots.clear();
out.extend(self.owners.drain().map(|(_, (obj, _))| obj));
out.extend(
self.groups
.drain()
.map(|(_, (group, _))| group.into_object()),
);
out.extend(self.snapshot.take().map(Into::into));
}
}

// Retention metadata is user-mutable. Its fast lookup must not call Python
// equality on a colliding key; non-exact string keys use the graph fallback.
struct ObjectKey<'a>(&'a str);
Expand Down Expand Up @@ -607,7 +681,7 @@ impl core::ops::Deref for CDataBytes {
pub(super) struct CDataValue {
pub bytes: CDataBytes,
keep: Option<PyObjectRef>,
hidden: Vec<PyObjectRef>,
hidden: Option<HiddenRefs>,
temporary: Option<PyObjectRef>,
}

Expand All @@ -624,7 +698,7 @@ impl CDataValue {
// Snapshot before writing: source and destination may overlap.
let bytes =
rustpython_host_env::ctypes::copy_to_sized_bytes(&source.buffer.read(), size);
let hidden = PyCData::hidden_refs(value);
let hidden = PyCData::hidden_refs(value, vm);
let keep = PyCData::get_kept_objects(value, vm);
return Ok(Self {
bytes: bytes.into(),
Expand Down Expand Up @@ -656,7 +730,7 @@ impl CDataValue {
.new_tuple(vec![PyCData::get_kept_objects(value, vm), value.to_owned()])
.into(),
),
hidden: Vec::new(),
hidden: None,
temporary: None,
});
}
Expand All @@ -682,7 +756,7 @@ impl CDataValue {
return Ok(Self {
bytes,
keep,
hidden: hidden.into_iter().collect(),
hidden: hidden.map(HiddenRefs::Single),
temporary: None,
});
}
Expand Down Expand Up @@ -716,7 +790,7 @@ impl CDataValue {
return Ok(Self {
bytes: vec![0; size].into(),
keep: None,
hidden: Vec::new(),
hidden: None,
temporary: None,
});
}
Expand All @@ -739,7 +813,7 @@ impl CDataValue {
let _temporary = self.temporary;
// Like KeepRef(None), preserve previous owners when assigning a raw
// address: it can still point into their private backing allocation.
if self.keep.is_none() && self.hidden.is_empty() {
if self.keep.is_none() && self.hidden.is_none() {
return Ok(());
}
// Publish every owner before dropping the displaced value, whose
Expand Down Expand Up @@ -1097,10 +1171,10 @@ impl PyCData {
/// Keep a reference alive without exposing it in _objects.
/// Retains a backing allocation used by a scalar's pointer.
pub fn keep_alive(&self, obj: PyObjectRef) {
drop(self.replace_hidden(0, vec![obj]));
drop(self.replace_hidden(0, Some(HiddenRefs::Single(obj))));
}

fn replace_hidden(&self, index: usize, owners: Vec<PyObjectRef>) -> Vec<PyObjectRef> {
fn replace_hidden(&self, index: usize, owners: Option<HiddenRefs>) -> Vec<PyObjectRef> {
let key = self.unique_key(index);
let base = self.base.read().clone();
let root = base.as_ref().map(|base| Self::find_root_object(base));
Expand All @@ -1109,7 +1183,7 @@ impl PyCData {
.and_then(|root| root.downcast_ref::<Self>())
.map_or(self, |cdata| cdata);
let mut refs = cdata.kept_refs.write();
if refs.is_none() && owners.is_empty() {
if refs.is_none() && owners.is_none() {
return Vec::new();
}
refs.get_or_insert_with(Default::default)
Expand Down Expand Up @@ -1194,13 +1268,35 @@ impl PyCData {

// Copying a view must retain its root's references, including private
// backing allocations and aliases through other fields or pointers.
// Keep the snapshot flat and deduplicate owners: self-copy must not nest
// a new holder around all previous snapshots.
fn hidden_refs(value: &PyObject) -> Vec<PyObjectRef> {
// Share a flat snapshot while the root's owner set is unchanged, including
// when copying between views of that root. Slots only retain group IDs.
fn hidden_refs(value: &PyObject, vm: &VirtualMachine) -> Option<HiddenRefs> {
let root = Self::find_root_object(value);
root.downcast_ref::<Self>()
.and_then(|cdata| cdata.kept_refs.read().as_ref().map(|refs| refs.snapshot()))
.unwrap_or_default()
let cdata = root.downcast_ref::<Self>()?;
let mut refs = cdata.kept_refs.write();
let refs = refs.as_mut()?;
match refs.owners.len() {
0 => None,
1 => Some(HiddenRefs::Single(
refs.owners.values().next().unwrap().0.clone(),
)),
_ => {
if refs.snapshot.is_none() {
// Snapshot publication must not call finalizers: raw bytes
// have already been copied, and owners must match those bytes.
// The tuple participates in GC immediately; only collection
// is left to the next ordinary allocation.
refs.snapshot = Some(PyRef::new_ref_without_collection(
PyTuple::new_unchecked(refs.snapshot().into_boxed_slice()),
vm.ctx.types.tuple_type.to_owned(),
None,
));
}
Some(HiddenRefs::Snapshot(
refs.snapshot.as_ref().unwrap().clone(),
))
}
}
}

// A scalar owns one reference; a scalar view stores it in its root instead.
Expand All @@ -1210,7 +1306,7 @@ impl PyCData {
hidden: Option<PyObjectRef>,
vm: &VirtualMachine,
) -> PyResult<()> {
let _old_hidden = self.replace_hidden(0, hidden.into_iter().collect());
let _old_hidden = self.replace_hidden(0, hidden.map(HiddenRefs::Single));
if self.base.read().is_some() {
self.keep_ref(0, value, vm)
} else {
Expand Down
Loading