Skip to content

Commit f197de0

Browse files
morealmoreal
committed
Avoid per-tail-call owner allocations
Assisted-by: Codex:GPT-5
1 parent eb08b18 commit f197de0

3 files changed

Lines changed: 53 additions & 47 deletions

File tree

‎crates/vm/src/frame.rs‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10680,11 +10680,10 @@ impl ExecutingFrame<'_> {
1068010680
}
1068110681
}
1068210682

10683-
// Pop the callable and transfer ownership to the trampoline via
10684-
// the VM side channel, avoiding a per-frame mutex lock on
10685-
// temporary_refs.
10683+
// Pop the callable and transfer ownership to the trampoline. This one
10684+
// reference keeps every field borrowed by the callee frame alive.
1068610685
let callable = self.pop_value();
10687-
unsafe { &mut *vm.pending_tailcall_refs.get() }.push(callable);
10686+
vm.set_pending_tailcall_owner(callable);
1068810687

1068910688
vm.set_pending_tailcall(callee_iframe);
1069010689
}
@@ -10734,13 +10733,13 @@ impl ExecutingFrame<'_> {
1073410733
*dst = Some(arg);
1073510734
}
1073610735
self.pop_value_opt(); // null (self_or_null)
10737-
let callable = self.pop_value(); // callable (bound method)
10736+
self.pop_value(); // callable (bound method)
1073810737
fastlocals[0] = Some(bound_self);
1073910738

10740-
// Transfer ownership to the trampoline via the VM side channel.
10741-
let refs = unsafe { &mut *vm.pending_tailcall_refs.get() };
10742-
refs.push(bound_function);
10743-
refs.push(callable);
10739+
// The function owns every field borrowed by the callee frame.
10740+
// bound_self is owned by fastlocals; the bound-method object itself is
10741+
// no longer needed and was dropped above, matching the recursive path.
10742+
vm.set_pending_tailcall_owner(bound_function);
1074410743

1074510744
vm.set_pending_tailcall(callee_iframe);
1074610745
}

‎crates/vm/src/vm/mod.rs‎

Lines changed: 44 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -110,11 +110,11 @@ pub struct VirtualMachine {
110110
/// pointer here before returning `ExecutionResult::TailCall`.
111111
/// Access only via `set_pending_tailcall` / `take_pending_tailcall`.
112112
pending_tailcall_frame: Cell<Option<PendingFrame>>,
113-
/// Owned references that keep callee raw pointers valid during TailCall.
114-
/// Set by `tailcall_prepare_frame`, drained by the trampoline into
115-
/// its local `owned_refs` Vec. Uses UnsafeCell because the VM is
116-
/// per-thread and this field is only accessed on the owning thread.
117-
pub(crate) pending_tailcall_refs: core::cell::UnsafeCell<Vec<PyObjectRef>>,
113+
/// Owned reference that keeps callee raw pointers valid during TailCall.
114+
/// Set by the exact-call handlers and moved into the trampoline's
115+
/// `SuspendedFrame`. Uses UnsafeCell because the VM is per-thread and this
116+
/// field is only accessed on the owning thread.
117+
pending_tailcall_owner: core::cell::UnsafeCell<Option<PyObjectRef>>,
118118
}
119119

120120
/// Non-owning frame pointer for the non-unix threading frames stack.
@@ -828,11 +828,11 @@ pub(crate) struct IframeEntryState {
828828
struct SuspendedFrame {
829829
iframe: *mut crate::frame::InterpreterFrame,
830830
entry_state: IframeEntryState,
831-
/// Owned references that keep callee's raw pointers (code, globals,
832-
/// builtins borrowed from PyFunction) valid. Drained from
833-
/// `vm.pending_tailcall_refs` when the callee's TailCall is consumed.
831+
/// Function that owns the callee's raw pointers (code, globals, builtins,
832+
/// closure, and func_obj). Moved from `vm.pending_tailcall_owner` when the
833+
/// callee's TailCall is consumed.
834834
/// Dropped when this SuspendedFrame is popped (after callee returns/errors).
835-
owned_refs: Vec<PyObjectRef>,
835+
callee_owner: PyObjectRef,
836836
/// True for the initial frame passed into the trampoline by the caller.
837837
/// The caller owns the datastack allocation for the entry frame, so the
838838
/// trampoline must NOT release it — only callee-allocated frames are
@@ -955,7 +955,7 @@ impl VirtualMachine {
955955
callable_cache: CallableCache::default(),
956956
audit_hooks: RefCell::new(vec![]),
957957
pending_tailcall_frame: Cell::new(None),
958-
pending_tailcall_refs: core::cell::UnsafeCell::new(Vec::with_capacity(2)),
958+
pending_tailcall_owner: core::cell::UnsafeCell::new(None),
959959
};
960960

961961
if vm.state.hash_secret.hash_str("")
@@ -1406,6 +1406,22 @@ impl VirtualMachine {
14061406
.set(Some(PendingFrame(core::ptr::NonNull::from(iframe))));
14071407
}
14081408

1409+
/// Store the function that owns the fields borrowed by the pending callee.
1410+
#[inline(always)]
1411+
pub(crate) fn set_pending_tailcall_owner(&self, owner: PyObjectRef) {
1412+
let slot = unsafe { &mut *self.pending_tailcall_owner.get() };
1413+
debug_assert!(slot.is_none(), "pending TailCall owner was not consumed");
1414+
*slot = Some(owner);
1415+
}
1416+
1417+
/// Take the pending callee owner, resetting the side channel.
1418+
#[inline(always)]
1419+
fn take_pending_tailcall_owner(&self) -> PyObjectRef {
1420+
unsafe { &mut *self.pending_tailcall_owner.get() }
1421+
.take()
1422+
.expect("TailCall without pending owner")
1423+
}
1424+
14091425
/// Take the pending tailcall frame pointer, resetting the side channel.
14101426
#[inline(always)]
14111427
fn take_pending_tailcall(&self) -> *mut crate::frame::InterpreterFrame {
@@ -1466,14 +1482,11 @@ impl VirtualMachine {
14661482
}
14671483

14681484
let initial_ptr = self.take_pending_tailcall();
1469-
// Drain the refs that keep the initial callee's raw pointers alive.
1470-
let initial_refs = unsafe { &mut *self.pending_tailcall_refs.get() }
1471-
.drain(..)
1472-
.collect();
1485+
let initial_owner = self.take_pending_tailcall_owner();
14731486
frame_stack.push(SuspendedFrame {
14741487
iframe: iframe as *mut crate::frame::InterpreterFrame,
14751488
entry_state,
1476-
owned_refs: initial_refs,
1489+
callee_owner: initial_owner,
14771490
is_entry: true,
14781491
});
14791492
let mut action = Action::EnterCallee(initial_ptr);
@@ -1498,13 +1511,11 @@ impl VirtualMachine {
14981511
let result = crate::frame::run_iframe(callee, self);
14991512
match result {
15001513
Ok(ExecutionResult::TailCall) => {
1501-
let refs = unsafe { &mut *self.pending_tailcall_refs.get() }
1502-
.drain(..)
1503-
.collect();
1514+
let callee_owner = self.take_pending_tailcall_owner();
15041515
frame_stack.push(SuspendedFrame {
15051516
iframe: callee_ptr,
15061517
entry_state: callee_entry,
1507-
owned_refs: refs,
1518+
callee_owner,
15081519
is_entry: false,
15091520
});
15101521
action = Action::EnterCallee(self.take_pending_tailcall());
@@ -1539,7 +1550,7 @@ impl VirtualMachine {
15391550
let SuspendedFrame {
15401551
iframe: caller_iframe_ptr,
15411552
entry_state: caller_entry,
1542-
owned_refs: _caller_refs,
1553+
callee_owner,
15431554
is_entry: caller_is_entry,
15441555
} = caller;
15451556
let caller_iframe = unsafe { &mut *caller_iframe_ptr };
@@ -1548,20 +1559,18 @@ impl VirtualMachine {
15481559
let result = crate::frame::run_iframe(caller_iframe, self);
15491560
match result {
15501561
Ok(ExecutionResult::TailCall) => {
1551-
let refs = unsafe { &mut *self.pending_tailcall_refs.get() }
1552-
.drain(..)
1553-
.collect();
1554-
drop(_caller_refs);
1562+
let next_callee_owner = self.take_pending_tailcall_owner();
1563+
drop(callee_owner);
15551564
frame_stack.push(SuspendedFrame {
15561565
iframe: caller_iframe_ptr,
15571566
entry_state: caller_entry,
1558-
owned_refs: refs,
1567+
callee_owner: next_callee_owner,
15591568
is_entry: caller_is_entry,
15601569
});
15611570
action = Action::EnterCallee(self.take_pending_tailcall());
15621571
}
15631572
Ok(ExecutionResult::Return(value)) => {
1564-
drop(_caller_refs);
1573+
drop(callee_owner);
15651574
self.exit_iframe(caller_entry);
15661575
if !caller_is_entry {
15671576
unsafe {
@@ -1574,7 +1583,7 @@ impl VirtualMachine {
15741583
}
15751584
Ok(ExecutionResult::Yield(_)) => panic!("Yield in non-generator frame"),
15761585
Err(exc) => {
1577-
drop(_caller_refs);
1586+
drop(callee_owner);
15781587
self.exit_iframe(caller_entry);
15791588
if !caller_is_entry {
15801589
unsafe {
@@ -1595,7 +1604,7 @@ impl VirtualMachine {
15951604
let SuspendedFrame {
15961605
iframe: caller_iframe_ptr,
15971606
entry_state: caller_entry,
1598-
owned_refs: _caller_refs,
1607+
callee_owner,
15991608
is_entry: caller_is_entry,
16001609
} = caller;
16011610
let caller_iframe = unsafe { &mut *caller_iframe_ptr };
@@ -1609,20 +1618,18 @@ impl VirtualMachine {
16091618
let result = crate::frame::run_iframe(caller_iframe, self);
16101619
match result {
16111620
Ok(ExecutionResult::TailCall) => {
1612-
let refs = unsafe { &mut *self.pending_tailcall_refs.get() }
1613-
.drain(..)
1614-
.collect();
1615-
drop(_caller_refs);
1621+
let next_callee_owner = self.take_pending_tailcall_owner();
1622+
drop(callee_owner);
16161623
frame_stack.push(SuspendedFrame {
16171624
iframe: caller_iframe_ptr,
16181625
entry_state: caller_entry,
1619-
owned_refs: refs,
1626+
callee_owner: next_callee_owner,
16201627
is_entry: caller_is_entry,
16211628
});
16221629
action = Action::EnterCallee(self.take_pending_tailcall());
16231630
}
16241631
Ok(ExecutionResult::Return(value)) => {
1625-
drop(_caller_refs);
1632+
drop(callee_owner);
16261633
self.exit_iframe(caller_entry);
16271634
if !caller_is_entry {
16281635
unsafe {
@@ -1639,7 +1646,7 @@ impl VirtualMachine {
16391646
panic!("Yield in non-generator frame")
16401647
}
16411648
Err(new_exc) => {
1642-
drop(_caller_refs);
1649+
drop(callee_owner);
16431650
self.exit_iframe(caller_entry);
16441651
if !caller_is_entry {
16451652
unsafe {
@@ -1655,7 +1662,7 @@ impl VirtualMachine {
16551662
}
16561663
}
16571664
Ok(Some(ExecutionResult::Return(value))) => {
1658-
drop(_caller_refs);
1665+
drop(callee_owner);
16591666
self.exit_iframe(caller_entry);
16601667
if !caller_is_entry {
16611668
unsafe {
@@ -1670,7 +1677,7 @@ impl VirtualMachine {
16701677
panic!("Unexpected execution result in trampoline unwind")
16711678
}
16721679
Err(new_exc) => {
1673-
drop(_caller_refs);
1680+
drop(callee_owner);
16741681
self.exit_iframe(caller_entry);
16751682
if !caller_is_entry {
16761683
unsafe {

‎crates/vm/src/vm/thread.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1023,7 +1023,7 @@ impl VirtualMachine {
10231023
callable_cache: self.callable_cache.clone(),
10241024
audit_hooks: RefCell::new(vec![]),
10251025
pending_tailcall_frame: Cell::new(None),
1026-
pending_tailcall_refs: core::cell::UnsafeCell::new(Vec::with_capacity(2)),
1026+
pending_tailcall_owner: core::cell::UnsafeCell::new(None),
10271027
};
10281028
ThreadedVirtualMachine { vm }
10291029
}

0 commit comments

Comments
 (0)