Skip to content

Optimize promise context removal #6267

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
proggeramlug merged 3 commits into from
Jul 11, 2026
Merged

Optimize promise context removal #6267

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
cc8e815
Optimize promise context removal
TheHypnoo Jul 10, 2026
a6c7312
Merge branch 'main' into perf/promise-context-removal
proggeramlug Jul 11, 2026
eaa2b01
Defer promise context rekeys during GC scan
TheHypnoo Jul 11, 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
245 changes: 223 additions & 22 deletions crates/perry-runtime/src/promise/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -657,34 +657,44 @@ pub(crate) fn enqueue_queue_microtask(callback: i64) {

#[derive(Default)]
pub(crate) struct PromiseContextStore {
entries: HashMap<usize, AsyncContextSnapshot>,
// The position stored with each snapshot makes removal and rekeying O(1)
// while `keys` remains the stable traversal surface for the GC scanners.
entries: HashMap<usize, (AsyncContextSnapshot, usize)>,
keys: Vec<usize>,
}

impl PromiseContextStore {
pub(crate) fn insert(&mut self, key: usize, snapshot: AsyncContextSnapshot) {
if !self.entries.contains_key(&key) {
self.keys.push(key);
if let Some((existing, _)) = self.entries.get_mut(&key) {
*existing = snapshot;
return;
}
self.entries.insert(key, snapshot);

let position = self.keys.len();
self.keys.push(key);
self.entries.insert(key, (snapshot, position));
}

pub(crate) fn get(&self, key: &usize) -> Option<&AsyncContextSnapshot> {
self.entries.get(key)
self.entries.get(key).map(|(snapshot, _)| snapshot)
}

pub(crate) fn get_mut(&mut self, key: &usize) -> Option<&mut AsyncContextSnapshot> {
self.entries.get_mut(key)
self.entries.get_mut(key).map(|(snapshot, _)| snapshot)
}

pub(crate) fn remove(&mut self, key: &usize) -> Option<AsyncContextSnapshot> {
let removed = self.entries.remove(key);
if removed.is_some() {
if let Some(pos) = self.keys.iter().position(|candidate| candidate == key) {
self.keys.swap_remove(pos);
}
let (snapshot, position) = self.entries.remove(key)?;
debug_assert_eq!(self.keys.get(position), Some(key));
let removed_key = self.keys.swap_remove(position);
debug_assert_eq!(removed_key, *key);
if let Some(moved_key) = self.keys.get(position) {
self.entries
.get_mut(moved_key)
.expect("PromiseContextStore key vector and map must agree")
.1 = position;
}
removed
Some(snapshot)
}

#[cfg(test)]
Expand All @@ -706,7 +716,7 @@ impl PromiseContextStore {
pub(crate) fn first(&self) -> Option<(usize, &AsyncContextSnapshot)> {
self.keys
.first()
.and_then(|key| self.entries.get(key).map(|snapshot| (*key, snapshot)))
.and_then(|key| self.get(key).map(|snapshot| (*key, snapshot)))
}

fn retain(&mut self, mut keep: impl FnMut(usize, &mut AsyncContextSnapshot) -> bool) {
Expand All @@ -716,12 +726,12 @@ impl PromiseContextStore {
let retain = self
.entries
.get_mut(&key)
.is_some_and(|snapshot| keep(key, snapshot));
.is_some_and(|(snapshot, _)| keep(key, snapshot));
if retain {
index += 1;
} else {
self.keys.swap_remove(index);
self.entries.remove(&key);
let removed = self.remove(&key);
debug_assert!(removed.is_some());
}
}
}
Expand All @@ -730,15 +740,206 @@ impl PromiseContextStore {
if old_key == new_key {
return;
}
let Some(context) = self.entries.remove(&old_key) else {
let Some((context, old_position)) = self.entries.remove(&old_key) else {
return;
};
if let Some(pos) = self.keys.iter().position(|key| *key == old_key) {
self.keys[pos] = new_key;
} else if !self.entries.contains_key(&new_key) {
self.keys.push(new_key);
debug_assert_eq!(self.keys.get(old_position), Some(&old_key));

if self.entries.contains_key(&new_key) {
// Preserve the moved promise's snapshot (the old key) and drop the
// stale target snapshot. The target key may be the last vector
// item, so update whichever key swap_remove relocates.
let removed_key = self.keys.swap_remove(old_position);
debug_assert_eq!(removed_key, old_key);
if let Some(moved_key) = self.keys.get(old_position) {
self.entries
.get_mut(moved_key)
.expect("PromiseContextStore key vector and map must agree")
.1 = old_position;
}
self.entries
.get_mut(&new_key)
.expect("rekey collision target must remain present")
.0 = context;
Comment on lines +743 to +763

@coderabbitai coderabbitai Bot Jul 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift

Prevent collision rekeying from skipping GC traversal entries.

scan_promise_contexts_step rekeys and then increments its index. On a collision, swap_remove(old_position) moves an unvisited key into that index, so it is skipped.

For example, rekeying 10 → 20 in [10, 20, 30] produces [30, 20]; the scanner processes 20, advances to index 1, processes 20 again, and never scans 30.

Defer collision rekeys until traversal completes, as the other scanner does, or explicitly coordinate index updates with the vector mutation. Extend the traversal regression test with a collision case.

Also applies to: 898-914

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/perry-runtime/src/promise/mod.rs` around lines 743 - 763, Collision
rekeying in the promise-context scanner mutates the keys vector and can skip an
unvisited entry. Update the relevant rekey logic in the promise context store
and scan_promise_contexts_step to defer collision rekeys until traversal
finishes, or adjust traversal indexing to revisit relocated entries safely;
apply the same fix to the corresponding logic around the alternate scanner and
extend the traversal regression test with a collision case.

} else {
self.keys[old_position] = new_key;
self.entries.insert(new_key, (context, old_position));
}
}

#[cfg(test)]
fn assert_invariants(&self) {
assert_eq!(self.entries.len(), self.keys.len());
for (position, key) in self.keys.iter().copied().enumerate() {
let (_, recorded_position) = self
.entries
.get(&key)
.expect("every key-vector entry must have a snapshot");
assert_eq!(*recorded_position, position);
}
self.entries.insert(new_key, context);
}
}

#[cfg(test)]
mod promise_context_store_bench {
use std::time::{Duration, Instant};

use super::{AsyncContextSnapshot, PromiseContextStore};

/// Reproducible regression probe for the former O(P²) key-vector search.
///
/// Run with:
/// `cargo test -p perry-runtime --release promise_context_store_remove_front \
/// -- --ignored --nocapture --test-threads=1`
#[test]
#[ignore = "manual performance regression probe"]
fn promise_context_store_remove_front() {
const BATCHES: [usize; 3] = [1_024, 4_096, 16_384];
const SAMPLES: usize = 5;

for batch in BATCHES {
let mut samples = Vec::with_capacity(SAMPLES);
for _ in 0..SAMPLES {
let mut store = PromiseContextStore::default();
for key in 0..batch {
store.insert(key, AsyncContextSnapshot::default());
}

let started = Instant::now();
for key in 0..batch {
std::hint::black_box(store.remove(&key));
}
samples.push(started.elapsed());
}
samples.sort_unstable();
let median = samples[SAMPLES / 2];
report(batch, median);
}
}

fn report(batch: usize, elapsed: Duration) {
println!(
"promise_context_store_remove_front batch={batch} median_ns={} ns_per_remove={}",
elapsed.as_nanos(),
elapsed.as_nanos() / batch as u128,
);
}
}

#[cfg(test)]
mod promise_context_store_tests {
use super::{AsyncContextSnapshot, PromiseContextStore};

fn snapshot(value: f64) -> AsyncContextSnapshot {
crate::async_context::test_snapshot_with_store(value)
}

fn stored_value(store: &PromiseContextStore, key: usize) -> Option<f64> {
store
.get(&key)
.and_then(crate::async_context::test_snapshot_first_store)
}

#[test]
fn duplicate_insert_missing_remove_and_last_swap_keep_index_consistent() {
let mut store = PromiseContextStore::default();
store.insert(10, snapshot(10.0));
store.insert(20, snapshot(20.0));
store.insert(30, snapshot(30.0));
store.insert(20, snapshot(200.0));

assert_eq!(store.keys, vec![10, 20, 30]);
assert_eq!(stored_value(&store, 20), Some(200.0));
assert!(store.remove(&99).is_none());
store.assert_invariants();

assert!(store.remove(&10).is_some());
assert_eq!(store.keys, vec![30, 20]);
assert_eq!(stored_value(&store, 30), Some(30.0));
store.assert_invariants();

assert!(store.remove(&20).is_some());
assert_eq!(store.keys, vec![30]);
store.assert_invariants();
}

#[test]
fn retain_partial_updates_positions_after_each_swap() {
let mut store = PromiseContextStore::default();
for key in 0..8 {
store.insert(key, snapshot(key as f64));
}

store.retain(|key, _| key % 2 == 0);

assert_eq!(store.entries.len(), 4);
for key in 0..8 {
assert_eq!(store.get(&key).is_some(), key % 2 == 0);
}
store.assert_invariants();
}

#[test]
fn rekey_collision_keeps_moved_snapshot_and_removes_duplicate_key() {
let mut store = PromiseContextStore::default();
store.insert(10, snapshot(10.0));
store.insert(20, snapshot(20.0));
store.insert(30, snapshot(30.0));

store.rekey(10, 20);

assert!(store.get(&10).is_none());
assert_eq!(stored_value(&store, 20), Some(10.0));
assert_eq!(stored_value(&store, 30), Some(30.0));
assert_eq!(store.keys.iter().filter(|&&key| key == 20).count(), 1);
store.assert_invariants();
}

#[test]
fn gc_relocation_rekeys_during_traversal_without_skipping_contexts() {
let mut store = PromiseContextStore::default();
for key in 0..4 {
store.insert(key, snapshot(key as f64));
}

let mut index = 0;
while let Some(old_key) = store.key_at(index) {
let new_key = old_key + 100;
store.rekey(old_key, new_key);
assert_eq!(stored_value(&store, new_key), Some(old_key as f64));
index += 1;
}

assert_eq!(store.keys, vec![100, 101, 102, 103]);
store.assert_invariants();
}

#[test]
fn deferred_rekey_collision_scans_every_context_before_swapping_keys() {
let mut store = PromiseContextStore::default();
store.insert(10, snapshot(10.0));
store.insert(20, snapshot(20.0));
store.insert(30, snapshot(30.0));

let mut scanned = Vec::new();
let mut moved = Vec::new();
let mut index = 0;
while let Some(key) = store.key_at(index) {
scanned.push(key);
if key == 10 {
moved.push((key, 20));
}
index += 1;
}
assert_eq!(scanned, vec![10, 20, 30]);

for (old_key, new_key) in moved {
store.rekey(old_key, new_key);
}

assert_eq!(stored_value(&store, 20), Some(10.0));
assert_eq!(stored_value(&store, 30), Some(30.0));
store.assert_invariants();
}
}

Expand Down
11 changes: 9 additions & 2 deletions crates/perry-runtime/src/promise/scanners.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,7 @@ pub(crate) struct PromiseRootScanState {
slot: usize,
context_entry: usize,
context_store: usize,
moved_context_keys: Vec<(usize, usize)>,
}

impl PromiseRootScanState {
Expand Down Expand Up @@ -517,12 +518,15 @@ fn scan_promise_contexts_step(
return false;
}
if visitor.visit_metadata_usize_slot(&mut new_key) {
contexts.rekey(key, new_key);
// Rekey only after every context has been scanned. A
// collision may swap an unvisited key into this slot;
// mutating the vector now would skip that context when
// the cursor advances.
state.moved_context_keys.push((key, new_key));
}
state.slot = 1;
}

let key = new_key;
let Some(context) = contexts.get_mut(&key) else {
state.index += 1;
state.finish_context_item();
Expand All @@ -540,6 +544,9 @@ fn scan_promise_contexts_step(
state.index += 1;
state.finish_context_item();
}
for (old_key, new_key) in state.moved_context_keys.drain(..) {
contexts.rekey(old_key, new_key);
}
true
})
}
Expand Down
Loading