Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughReplaces epoch-based GC with a permit-based coordinator and implements a generational collector (Gen0/Gen1/Gen2 + inactive) with a card table, VM early-yield and write-barrier integration, boxes several Object payloads, and updates engine/heap/VM APIs and tests. Changes
Sequence Diagram(s)sequenceDiagram
participant App as Application
participant Engine as BexEngine
participant PermitMgr as HeapPermitManager
participant HeapGuard as HeapGuard
participant Permit as ActiveHeapPermit
participant VM as BexVm
participant Heap as BexHeap
rect rgba(100, 150, 200, 0.5)
App->>Engine: call_function(...)
Engine->>PermitMgr: new_permit(...) -> InactiveHeapPermit
Engine->>Permit: acquire() -> ActiveHeapPermit
Permit->>VM: Deref -> access VM & heap
VM->>Heap: allocate / mutate (write barrier marks cards)
end
rect rgba(200, 100, 100, 0.5)
App->>Engine: collect_garbage(level)
Engine->>PermitMgr: request_park()
PermitMgr->>HeapGuard: await active permits (drain)
HeapGuard->>Heap: collect_garbage_generational(roots, level)
Heap->>Engine: return (GcStats, forwarding_map)
Engine->>HeapGuard: forward_roots(forwarding_map)
end
rect rgba(100, 200, 100, 0.5)
VM->>VM: EarlyYieldCheck triggers -> VmExecState::EarlyYield
VM->>Permit: release/renew to cooperate with GC
VM->>VM: resume execution after permit reacquire
end
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Binary size checks passed✅ 7 passed
Generated by |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (5)
baml_language/crates/bex_heap/src/heap_debugger/real.rs (1)
238-248:⚠️ Potential issue | 🟠 MajorInclude Gen1 and Gen2 in full heap verification.
verify_full_impl()now walks only Gen0, so promoted objects in Gen1/Gen2 can violate invariants without being checked underHeapVerifyMode::Full.Suggested fix
- // Verify all runtime objects in Gen0 (the active nursery) + // Verify all runtime objects across all generations. unsafe { - let gen0 = &*self.gen0.get(); - for (runtime_idx, obj) in gen0.iter().enumerate() { - let ptr = gen0.get_ptr(runtime_idx); - let idx = HeapPtr::from_ptr(ptr, self.heap_epoch()); - if self.debug_handle_runtime_sentinel(idx, obj, ct_len) { - continue; - } - self.verify_object_invariants(idx, obj, ct_len); + for generation in [&*self.gen0.get(), &*self.gen1.get(), &*self.gen2.get()] { + for (runtime_idx, obj) in generation.iter().enumerate() { + let ptr = generation.get_ptr(runtime_idx); + let idx = HeapPtr::from_ptr(ptr, self.heap_epoch()); + if self.debug_handle_runtime_sentinel(idx, obj, ct_len) { + continue; + } + self.verify_object_invariants(idx, obj, ct_len); + } } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@baml_language/crates/bex_heap/src/heap_debugger/real.rs` around lines 238 - 248, The full-heap verifier (verify_full_impl) currently iterates only the active nursery (gen0), so promoted objects in gen1 and gen2 are skipped under HeapVerifyMode::Full; update verify_full_impl to also iterate gen1 and gen2 similarly to gen0: obtain references to self.gen1 and self.gen2, enumerate their entries, compute ptr and idx via gen?.get_ptr and HeapPtr::from_ptr with self.heap_epoch(), call debug_handle_runtime_sentinel(idx, obj, ct_len) and then verify_object_invariants(idx, obj, ct_len) when not handled, mirroring the same loop logic used for gen0 so all generations are checked under Full verification.baml_language/crates/bex_vm/src/watch.rs (1)
558-591:⚠️ Potential issue | 🟠 MajorCollect every
RootStatepointer thatforward_rootscan rewrite.
forward_roots()updatesstate.valueandWatchFilter::Function, butcollect_roots()does not report them. If either points to a movable runtime object, GC can collect or fail to forward it before this state is patched.Suggested fix
fn collect_roots(&self, roots: &mut Vec<HeapPtr>) { for state in self.roots.values() { + if let Value::Object(ptr) = state.value { + roots.push(ptr); + } if let Some(Value::Object(ptr)) = state.last_assigned { roots.push(ptr); } if let Some(Value::Object(ptr)) = state.last_notified { roots.push(ptr); } + if let WatchFilter::Function(ptr) = state.filter { + roots.push(ptr); + } } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@baml_language/crates/bex_vm/src/watch.rs` around lines 558 - 591, collect_roots currently only pushes last_assigned/last_notified but forward_roots also forwards state.value and WatchFilter::Function, so update collect_roots to report every HeapPtr that forward_roots can rewrite: inspect state.value and if it is Value::Object push that ptr, and if state.filter is WatchFilter::Function push that function ptr as well (in addition to the existing checks for last_assigned and last_notified). This ensures collect_roots, forward_roots and the RootState/WatchFilter handling are synchronized.baml_language/crates/bex_engine/src/lib.rs (2)
794-814:⚠️ Potential issue | 🟠 MajorRegister active calls before execution.
cancel_function_call()can only cancel entries inactive_calls, butcall_function()never insertscall_id. This makes ID-based cancellation returnFunctionCallNotFoundfor in-flight calls and leavesDuplicateCallIdeffectively unused.🛠️ Possible fix shape
+struct ActiveCallGuard<'a> { + engine: &'a BexEngine, + call_id: CallId, +} + +impl Drop for ActiveCallGuard<'_> { + fn drop(&mut self) { + self.engine.active_calls.lock().unwrap().remove(&self.call_id); + } +} + pub async fn call_function( self: &Arc<Self>, function_name: &str, @@ if cancel.is_cancelled() { return Err(EngineError::Cancelled); } + + { + let mut active_calls = self.active_calls.lock().unwrap(); + if active_calls.contains_key(&call_id) { + return Err(EngineError::DuplicateCallId { call_id }); + } + active_calls.insert(call_id, cancel.clone()); + } + let _active_call = ActiveCallGuard { + engine: self, + call_id, + };- if let Some(cancel) = active_calls.remove(&call_id) { + if let Some(cancel) = active_calls.get(&call_id) { cancel.cancel(); Ok(())Also applies to: 923-930
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@baml_language/crates/bex_engine/src/lib.rs` around lines 794 - 814, call_function currently never registers the call_id into active_calls so cancel_function_call cannot find in-flight calls; modify call_function to validate and insert the call_id into the active_calls map before creating the VM/executing (use the same active_calls structure manipulated by cancel_function_call), returning DuplicateCallId if the id already exists, and ensure you remove the entry on function completion or error (and similarly apply the same registration/removal logic in the other call path referenced around the later block). Target the functions/methods call_function, cancel_function_call and the active_calls data structure to add the insert/check (for duplicate), and ensure cleanup on all exit paths.
1389-1433:⚠️ Potential issue | 🔴 CriticalRelease the heap permit while waiting on async futures.
This
Awaitpath callsgc_safepoint()but then waits onprocessed_futures.recv()while still holdingActiveHeapPermit<BexVm>. A long-running sys-op can block GC for its full duration; worse, if GC is waiting and the sys-op tries to spawn another VM, it can deadlock behind the GC-held holders mutex while this VM still holds the permit. Release to an inactive permit before thetokio::select!, then reacquire only to write completed results back into the VM.🛠️ Possible fix shape
VmExecState::Await(future_id) => { Self::cancellation_safepoint(cancel, &abort_handles)?; vm = self.gc_safepoint(vm).await; @@ - loop { + let mut inactive_vm = vm.release(); + loop { tokio::select! { biased; () = cancel.cancelled() => { // Abort all in-flight spawned tasks to stop // HTTP requests, sleeps, etc. immediately. abort_handles.abort_all(); return Err(EngineError::Cancelled); } future = processed_futures.recv() => { let future = future .ok_or(EngineError::FutureChannelClosed)?; let external = future.result?; + vm = inactive_vm.acquire().await; let value = self.convert_external_to_vm_value( &mut vm, external, ); vm.fulfil_future(future.id, value) .map_err(EngineError::VmInternalError)?; if future.id == future_id { break; } + inactive_vm = vm.release(); } } } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@baml_language/crates/bex_engine/src/lib.rs` around lines 1389 - 1433, The Await branch currently holds an ActiveHeapPermit<BexVm> across the tokio::select!, blocking GC; change VmExecState::Await to release the active permit returned by gc_safepoint() (or convert it to an inactive permit) before entering the select/loop that awaits processed_futures.recv(), then reacquire or upgrade back to an active permit only when you need to call vm.fulfil_future/convert_external_to_vm_value; ensure cancellation_safepoint(cancel, &abort_handles)? still runs before releasing and that abort_handles.abort_all() is used on cancel, and update calls around processed_futures.try_recv()/recv() to perform VM writes only while the active permit is held.baml_language/crates/bex_vm/src/vm.rs (1)
2186-2191:⚠️ Potential issue | 🔴 CriticalInclude VM globals in GC root collection and forwarding.
StoreGlobalcan retain a runtimeValue::Objectinself.globals, butRootHaver for BexVmonly scans stack/watch/frames. A live object reachable only from a global can be collected during GC, and moved global pointers won’t be forwarded.🐛 Proposed fix
impl ::bex_vm_types::RootHaver for BexVm { fn collect_roots(&self, roots: &mut Vec<HeapPtr>) { // Stack values roots.extend(self.stack.iter().filter_map(|v| match v { Value::Object(ptr) => Some(*ptr), _ => None, })); + // Global values + roots.extend(self.globals.iter().filter_map(|v| match v { + Value::Object(ptr) => Some(*ptr), + _ => None, + })); + // Watch state (last_assigned/last_notified values that aren't on the stack) self.watch.collect_roots(roots); // Frame function pointers (needed once closures are heap-allocated) roots.extend(self.collect_frame_roots()); @@ for value in &mut self.stack { if let Value::Object(ptr) = value { if let Some(&new_ptr) = roots.get(ptr) { *ptr = new_ptr; } } } + // Global values + for value in &mut self.globals { + if let Value::Object(ptr) = value { + if let Some(&new_ptr) = roots.get(ptr) { + *ptr = new_ptr; + } + } + } + // Watch state (last_assigned/last_notified values that aren't on the stack) self.watch.forward_roots(roots);Adapt the iterator syntax to
GlobalPool’s actual API if needed.Also applies to: 3905-3947
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@baml_language/crates/bex_vm/src/vm.rs` around lines 2186 - 2191, The VM currently stores objects into self.globals via Instruction::StoreGlobal but RootHaver for BexVm does not include globals in GC root scanning/forwarding, so globals can be collected or not updated after moving; update the BexVm RootHaver implementation to iterate over the VM globals (self.globals / GlobalPool) and for each entry call the same root-scanning/forwarding helper used for stack/frames (e.g., scan_root/forward_root or the existing RootHaver APIs), adapting the iteration to GlobalPool’s actual iterator/method names, and ensure both collection and pointer-forwarding paths treat globals the same as stack/watch/frames so Value::Object in globals are preserved and updated during GC.
🧹 Nitpick comments (6)
baml_language/crates/bex_engine/tests/early_yield.rs (1)
56-64: Use a deterministic rendezvous before requesting GC.
yield_now()only gives the scheduler opportunities; it does not guarantee the spawned VM has entered the loop or acquired a heap permit. A test hook/barrier inside the call path would prevent these tests from passing when GC happens before execution is actually in flight.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@baml_language/crates/bex_engine/tests/early_yield.rs` around lines 56 - 64, The test currently uses let_call_start() which repeatedly tokio::task::yield_now() to hope the spawned call_function enters its hot loop; replace this with a deterministic rendezvous by adding and using a test-only synchronization point inside the VM call path (e.g., a barrier/Notify/oneshot signal) that the call_function will signal when it has acquired its heap permit/entered the hot loop, and have the test wait on that signal before requesting GC; update the test to wait on that new hook instead of calling let_call_start(), and ensure the hook is only enabled in test builds or behind a feature so production code is unaffected.baml_language/crates/bex_heap/src/gc.rs (5)
710-718: Inconsistency:debug_verify_tlab_canaries()is missing from the minor path.
copy_collectioncallsself.debug_verify_tlab_canaries()at L136 before doing anything.collect_garbage_minorskips it. Since minor GC will be by far the more frequent collection, it's the more valuable place for the canary check. Add the same call at the top ofcollect_garbage_minorfor parity.♻️ Proposed diff
pub unsafe fn collect_garbage_minor( &self, roots: &[HeapPtr], ) -> (GcStats, Vec<HeapPtr>, HashMap<HeapPtr, HeapPtr>) { let mut forwarding: HashMap<HeapPtr, HeapPtr> = HashMap::new(); + self.debug_verify_tlab_canaries(); + self.bump_epoch();🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@baml_language/crates/bex_heap/src/gc.rs` around lines 710 - 718, The minor-GC path is missing the TLAB canary verification present in the full copy collection; add a call to self.debug_verify_tlab_canaries() at the start of collect_garbage_minor so it performs the same canary check as copy_collection (which already invokes debug_verify_tlab_canaries()). Place the call before any mutation or epoch bump (i.e., at the top of collect_garbage_minor, analogous to the call at the start of copy_collection) to ensure parity and early detection of TLAB corruption.
867-876: Minor:collect_garbage_generationalcan be#[inline]and letMajorskip the trampoline.Tiny dispatcher that just forwards to one of two functions. Either inline it or have callers pick the specific entry point directly. Not important — noting it only because both arms are already exposed as
pub unsafe fnonBexHeap, so the dispatcher is more of a convenience than a necessity.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@baml_language/crates/bex_heap/src/gc.rs` around lines 867 - 876, Add #[inline] to the collect_garbage_generational method to allow the trivial dispatcher to be inlined, or remove the method and update callers to invoke collect_garbage_minor for CollectionLevel::Minor and collect_garbage for CollectionLevel::Major directly; specifically modify the BexHeap method collect_garbage_generational (which matches on CollectionLevel::Minor / CollectionLevel::Major) so it either has the #[inline] attribute or eliminate its usage and adjust call sites to call collect_garbage_minor or collect_garbage directly.
2434-2525: These two "tests" don't assert the thing they claim to test.
test_should_collect_minor_on_gen1_pressure(L2436-2477) andtest_should_collect_major_on_gen2_pressure(L2481-2525) are mostly comment blocks explaining why the scenario can't be triggered without a test-only API. The only real assertions are either a Gen0-pressureMinortrigger (already covered bytest_should_collect_minor_on_gen0_pressure) or aNoneobservation. Neither test actually exercises the Gen1-threshold branch (L78) or the Gen2-threshold branch (L71) ofshould_collect.
test_should_collect_minor_when_gen1_exceeds_adapted_thresholddoes cover the Gen1 path, so the gap is specifically the Gen2→Majortrigger. Suggest either:
- deleting the two non-tests (the misleading coverage is worse than no test), or
- adding a minimal
#[cfg(test)] pub(crate)setter likeset_gen2_collection_threshold_for_test(usize)so both branches can be unit-tested directly without the 50k-object allocation.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@baml_language/crates/bex_heap/src/gc.rs` around lines 2434 - 2525, The two tests (test_should_collect_minor_on_gen1_pressure and test_should_collect_major_on_gen2_pressure) do not actually exercise the Gen1/Gen2 threshold branches of should_collect; either remove these misleading no-op tests or add a test-only setter so the thresholds can be driven directly. Implement a #[cfg(test)] pub(crate) setter (e.g., set_gen1_collection_threshold_for_test(usize) and set_gen2_collection_threshold_for_test(usize)) in the heap module, use those in the tests to set small thresholds, then allocate/promote a few objects and assert should_collect() returns Some(CollectionLevel::Minor) or Some(CollectionLevel::Major); reference helpers update_thresholds_after_minor/update_thresholds_after_major and the should_collect function when updating tests.
739-768: Document (or debug-assert) the card-table invariant theGen2arm depends on.Roots handed to
collect_garbage_minorare not filtered by generation, so a Gen2 pointer can land directly in the worklist. TheGeneration::Gen2arm identity-maps it and does not trace its outgoing references. Correctness here rests entirely on the invariant that every Gen2→young reference is covered by a dirty card (major GC produces only Gen2/compile-time refs in Gen2; the promotion sweep at L805-814 marks promoted objects dirty; mutator write barriers mark on store).If any future code path manages to create a Gen2→young edge without dirtying the card (e.g. a new direct-Gen2 allocation path, a missed write barrier, or a bulk copy that bypasses the barrier), this arm will silently drop live young objects.
Worth adding a short comment making that invariant explicit and, ideally, a
debug_assert!in debug builds that walks this Gen2 root's references and asserts either!is_young(child)or thatchildis already inforwarding/ its card is dirty. Cheap insurance against future regressions.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@baml_language/crates/bex_heap/src/gc.rs` around lines 739 - 768, The Gen2 arm currently identity-maps Gen2 roots without tracing their outgoing references, relying on the card-table invariant; add a short comment in the Generation::Gen2 branch documenting that every Gen2→young edge must be covered by a dirty card (maintained by major GC, promotion sweep and the mutator write barrier), and add a debug-only assertion that for each outgoing reference of the Gen2 root (use the same path as add_references_to_worklist to iterate children) either the child is not young (use self.generation_of(child) != Generation::Gen0), or the child is already in forwarding, or the corresponding card is dirty (call your existing card-check helper / card table API); place this debug_assert before forwarding.insert(old_ptr, old_ptr) to catch invariant violations early.
805-825: Dirty cards are never cleared across minor GCs — card set grows monotonically.Major GC clears and resizes
gen2_cards(L204-208). Minor GC only adds marks (dirty-card scan + promotion sweep) and never clears. Consequences:
- A card that became dirty because of a write to a Gen0 object later promoted to Gen2 and then overwritten with a primitive stays dirty forever.
- The promotion write-barrier sweep marks every newly-promoted Gen2 object regardless of whether it actually holds a young reference, so each minor GC monotonically increases the dirty set.
- Over many minor GCs the dirty-card scan degenerates toward a full Gen2 scan, nullifying the generational optimization.
Standard fix is to clear dirty cards at the end of the minor GC (after fixup + promotion sweep) and re-dirty from the promotion sweep only — i.e. clear before L805, since by that point all stale cards have been fixed up, and the sweep itself will restore the marks for promoted objects. Other cards can be conservatively kept dirty only if you can prove the card still has young refs after fixup, or simply cleared if the design intends the write barrier to be the sole source of truth post-GC.
♻️ Sketch of where the reset fits
+ // Clear the Gen2 card table — after fixup, all pre-existing dirty + // cards have either had their young refs updated (to still-young + // targets, which will be re-dirtied below if promoted, or are + // genuinely still young and will be re-discovered by future write + // barriers on the next mutator store) or no longer contain young + // refs at all. The promotion sweep below re-establishes the + // remembered set for freshly-promoted objects. + // SAFETY: GC safepoint. + unsafe { (*self.gen2_cards.get()).clear(); } + // Promotion write-barrier: mark the card dirty for every newly-promoted // Gen2 object. ...Note: this only works if pre-existing Gen2 objects with surviving young refs can be re-dirtied via the mutator's write barrier on the next store. If you need the remembered set to remain complete across minor GCs without a store, the clear needs to be replaced by a targeted "re-dirty if still points young" sweep of the existing dirty cards after fixup.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@baml_language/crates/bex_heap/src/gc.rs` around lines 805 - 825, The gen2 dirty-card set is never cleared during minor GC so it grows monotonically; fix it by clearing the remembered-set right after the minor-GC fixup/promotion sweep and before re-dirtying promoted objects — e.g. after promotion/fixup and before the loop that calls gen2_ref() + mark_card_for_ptr(), call the gen2_cards clear/reset method (the same structure manipulated by ensure_capacity_for_chunks and cleared during major GC) so that stale cards are removed and only the promotion sweep/write-barrier re-dirties needed slots; touch the code around gen2_ref, gen2_cards, mark_card_for_ptr, and finalize_inactive_space to insert the clear at that point.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@baml_language/crates/bex_engine/src/heap_guard.rs`:
- Around line 25-68: ActiveHeapPermit<T> can auto-implement Sync for non-Sync
root holders, allowing concurrent access to a non-Sync holder; add a
compile-time test that verifies ActiveHeapPermit<NonSyncRootHaver> is Send but
not Sync by creating a small non-Sync type (e.g., struct NonSyncRootHaver { /*
!Sync */ }), then in a test or compile-time assertions use static_assertions
macros (assert_impl_all! and assert_not_impl_any!) to assert
ActiveHeapPermit<NonSyncRootHaver>: Send and not Sync; if the test fails, make
ActiveHeapPermit explicitly non-Sync by adding a PhantomData<*mut T> or other
marker to the ActiveHeapPermit struct to prevent automatic Sync impls (adjust
the struct definition and rerun the assertions).
In `@baml_language/crates/bex_engine/tests/registry.rs`:
- Around line 104-108: The test currently only checks the enum variant after GC,
which can be a false positive; instead bind the handle (the
BexExternalValue::Handle variant from result) to a local variable, call the
engine API that dereferences/reads a handle (e.g., engine.resolve_handle or
engine.dereference_handle — use the actual method your engine exposes to read a
handle) after collect_garbage(CollectionLevel::Major). Then assert the
dereferenced value is the expected forwarded object (or that dereference
succeeds), replacing the matches!(result, BexExternalValue::Handle(_)) check so
the test proves the slab key still points to a valid forwarded object.
In `@baml_language/crates/bex_heap/src/card_table.rs`:
- Around line 85-94: The current mark_dirty_by_offset silently no-ops when
card_index is out of bounds; change it to make out-of-bounds marks observable by
either returning a bool/Result or at minimum adding a debug assertion. Update
the signature of mark_dirty_by_offset (or overload) to return a bool indicating
success/failure, or add a debug_assert!(card_index < self.cards.len(), "...")
before the get, and if returning a bool, return false when
self.cards.get(card_index) is None and true after slot.store(1,
Ordering::Relaxed). Reference mark_dirty_by_offset, CARDS_PER_CHUNK,
ensure_capacity_for_chunks and the self.cards.get(...) / slot.store(...) calls
so reviewers can locate and adjust callers to handle the new return value or
rely on the debug_assert in debug builds.
In `@baml_language/crates/bex_heap/src/chunked_vec.rs`:
- Around line 194-218: The two public helpers num_chunks and chunk_start_ptr
must be made unsafe to reflect their safepoint/bounds preconditions: change
their signatures to `pub unsafe fn num_chunks(&self) -> usize` and `pub unsafe
fn chunk_start_ptr(&self, chunk_idx: usize) -> *const T`, keep the existing
SAFETY doc comments (and update the top-level doc to state callers must call
them inside unsafe) and then update any call sites to invoke them inside unsafe
blocks so callers explicitly acknowledge the no-concurrent-growth and bounds
requirements.
In `@baml_language/crates/bex_heap/src/gc.rs`:
- Line 836: collect_garbage_minor currently builds a partial forwarding map
(only visited objects) and then calls update_handles which invalidates any
handle whose target isn't in forwarding, causing live Gen2 handles to be
dropped; to fix, either (A) before calling update_handles in
collect_garbage_minor, walk the handle roots produced by collect_handle_roots
and insert identity forwarding entries for any live Gen2 targets into the
forwarding map (so Gen2 handles have entries), or (B) change update_handles to
be minor-GC-aware (when invoked from collect_garbage_minor) and treat missing
forwarding entries as "keep handle as-is" instead of invalidating (i.e.,
forward-if-present, else keep); update the call site and add a clear
conditional/path in collect_garbage_minor or update_handles to select the
correct behavior.
In `@baml_language/crates/bex_heap/src/heap.rs`:
- Around line 642-647: The len() and stats() implementations only include Gen0
so promoted objects in Gen1/Gen2 are omitted; update len() to sum live counts
from all generations by reading (*self.gen0.get()).len(),
(*self.gen1.get()).len(), and (*self.gen2.get()).len() (safely using the same
pattern as gen0) and update stats() (and the HeapStats::runtime_objects
calculation) to aggregate runtime object counts across gen0, gen1, and gen2
instead of just gen0 so promotions are reflected in heap length and stats.
In `@baml_language/crates/bex_heap/src/tlab.rs`:
- Around line 253-259: Update the doc comment in tlab.rs under the "Write
Barrier" section to correctly state the barrier contract: explain that a write
barrier must be fired when an older-generation object (i.e., the target object
being mutated) is written with a reference to a younger-generation object
(old-to-young), not when a newly allocated object points to older targets.
Reference the "Write Barrier" heading in this file and adjust the wording so
callers know to invoke the barrier when mutating older-generation objects that
receive younger HeapPtr references.
In `@baml_language/crates/bex_heap/tests/generational.rs`:
- Around line 528-546: The test currently marks a card that only contains Null
so it won't keep orphan2 alive even if the card table isn't cleared; change the
test to make the dirty card actually reference a live Gen0 object before the
first GC (e.g., allocate a Gen0 object via tlab.alloc_string or similar, store
its pointer into the Gen2 slot you mark with mark_card_for_ptr so the dirty card
truly reflects a cross-generational pointer), then run the first minor GC and
verify that after the second minor GC orphan2 is collected only if the dirty
card was cleared; alternatively (or additionally) use/add a test-only
heap/card-table introspection API to assert the dirty-card count before/after
the first collect_garbage_generational call to ensure the card was cleared.
Ensure you update references to mark_card_for_ptr, collect_garbage_generational,
tlab.alloc_string, orphan2 and stats.collected_count in the test.
In `@baml_language/crates/tools_onionskin/src/compiler.rs`:
- Around line 3850-3853: BexVm::from_program is being called with park_requested
unconditionally but that parameter is only available on non-wasm builds; guard
creation and passing of park_requested with cfg(not(target_arch = "wasm32")) and
call BexVm::from_program without the extra argument on wasm targets.
Specifically, wrap the Arc::new(...) and the call that includes park_requested
in a #[cfg(not(target_arch = "wasm32"))] block, and provide an alternative
#[cfg(target_arch = "wasm32")] branch that calls BexVm::from_program(program)
with no park_requested argument.
---
Outside diff comments:
In `@baml_language/crates/bex_engine/src/lib.rs`:
- Around line 794-814: call_function currently never registers the call_id into
active_calls so cancel_function_call cannot find in-flight calls; modify
call_function to validate and insert the call_id into the active_calls map
before creating the VM/executing (use the same active_calls structure
manipulated by cancel_function_call), returning DuplicateCallId if the id
already exists, and ensure you remove the entry on function completion or error
(and similarly apply the same registration/removal logic in the other call path
referenced around the later block). Target the functions/methods call_function,
cancel_function_call and the active_calls data structure to add the insert/check
(for duplicate), and ensure cleanup on all exit paths.
- Around line 1389-1433: The Await branch currently holds an
ActiveHeapPermit<BexVm> across the tokio::select!, blocking GC; change
VmExecState::Await to release the active permit returned by gc_safepoint() (or
convert it to an inactive permit) before entering the select/loop that awaits
processed_futures.recv(), then reacquire or upgrade back to an active permit
only when you need to call vm.fulfil_future/convert_external_to_vm_value; ensure
cancellation_safepoint(cancel, &abort_handles)? still runs before releasing and
that abort_handles.abort_all() is used on cancel, and update calls around
processed_futures.try_recv()/recv() to perform VM writes only while the active
permit is held.
In `@baml_language/crates/bex_heap/src/heap_debugger/real.rs`:
- Around line 238-248: The full-heap verifier (verify_full_impl) currently
iterates only the active nursery (gen0), so promoted objects in gen1 and gen2
are skipped under HeapVerifyMode::Full; update verify_full_impl to also iterate
gen1 and gen2 similarly to gen0: obtain references to self.gen1 and self.gen2,
enumerate their entries, compute ptr and idx via gen?.get_ptr and
HeapPtr::from_ptr with self.heap_epoch(), call
debug_handle_runtime_sentinel(idx, obj, ct_len) and then
verify_object_invariants(idx, obj, ct_len) when not handled, mirroring the same
loop logic used for gen0 so all generations are checked under Full verification.
In `@baml_language/crates/bex_vm/src/vm.rs`:
- Around line 2186-2191: The VM currently stores objects into self.globals via
Instruction::StoreGlobal but RootHaver for BexVm does not include globals in GC
root scanning/forwarding, so globals can be collected or not updated after
moving; update the BexVm RootHaver implementation to iterate over the VM globals
(self.globals / GlobalPool) and for each entry call the same
root-scanning/forwarding helper used for stack/frames (e.g.,
scan_root/forward_root or the existing RootHaver APIs), adapting the iteration
to GlobalPool’s actual iterator/method names, and ensure both collection and
pointer-forwarding paths treat globals the same as stack/watch/frames so
Value::Object in globals are preserved and updated during GC.
In `@baml_language/crates/bex_vm/src/watch.rs`:
- Around line 558-591: collect_roots currently only pushes
last_assigned/last_notified but forward_roots also forwards state.value and
WatchFilter::Function, so update collect_roots to report every HeapPtr that
forward_roots can rewrite: inspect state.value and if it is Value::Object push
that ptr, and if state.filter is WatchFilter::Function push that function ptr as
well (in addition to the existing checks for last_assigned and last_notified).
This ensures collect_roots, forward_roots and the RootState/WatchFilter handling
are synchronized.
---
Nitpick comments:
In `@baml_language/crates/bex_engine/tests/early_yield.rs`:
- Around line 56-64: The test currently uses let_call_start() which repeatedly
tokio::task::yield_now() to hope the spawned call_function enters its hot loop;
replace this with a deterministic rendezvous by adding and using a test-only
synchronization point inside the VM call path (e.g., a barrier/Notify/oneshot
signal) that the call_function will signal when it has acquired its heap
permit/entered the hot loop, and have the test wait on that signal before
requesting GC; update the test to wait on that new hook instead of calling
let_call_start(), and ensure the hook is only enabled in test builds or behind a
feature so production code is unaffected.
In `@baml_language/crates/bex_heap/src/gc.rs`:
- Around line 710-718: The minor-GC path is missing the TLAB canary verification
present in the full copy collection; add a call to
self.debug_verify_tlab_canaries() at the start of collect_garbage_minor so it
performs the same canary check as copy_collection (which already invokes
debug_verify_tlab_canaries()). Place the call before any mutation or epoch bump
(i.e., at the top of collect_garbage_minor, analogous to the call at the start
of copy_collection) to ensure parity and early detection of TLAB corruption.
- Around line 867-876: Add #[inline] to the collect_garbage_generational method
to allow the trivial dispatcher to be inlined, or remove the method and update
callers to invoke collect_garbage_minor for CollectionLevel::Minor and
collect_garbage for CollectionLevel::Major directly; specifically modify the
BexHeap method collect_garbage_generational (which matches on
CollectionLevel::Minor / CollectionLevel::Major) so it either has the #[inline]
attribute or eliminate its usage and adjust call sites to call
collect_garbage_minor or collect_garbage directly.
- Around line 2434-2525: The two tests
(test_should_collect_minor_on_gen1_pressure and
test_should_collect_major_on_gen2_pressure) do not actually exercise the
Gen1/Gen2 threshold branches of should_collect; either remove these misleading
no-op tests or add a test-only setter so the thresholds can be driven directly.
Implement a #[cfg(test)] pub(crate) setter (e.g.,
set_gen1_collection_threshold_for_test(usize) and
set_gen2_collection_threshold_for_test(usize)) in the heap module, use those in
the tests to set small thresholds, then allocate/promote a few objects and
assert should_collect() returns Some(CollectionLevel::Minor) or
Some(CollectionLevel::Major); reference helpers
update_thresholds_after_minor/update_thresholds_after_major and the
should_collect function when updating tests.
- Around line 739-768: The Gen2 arm currently identity-maps Gen2 roots without
tracing their outgoing references, relying on the card-table invariant; add a
short comment in the Generation::Gen2 branch documenting that every Gen2→young
edge must be covered by a dirty card (maintained by major GC, promotion sweep
and the mutator write barrier), and add a debug-only assertion that for each
outgoing reference of the Gen2 root (use the same path as
add_references_to_worklist to iterate children) either the child is not young
(use self.generation_of(child) != Generation::Gen0), or the child is already in
forwarding, or the corresponding card is dirty (call your existing card-check
helper / card table API); place this debug_assert before
forwarding.insert(old_ptr, old_ptr) to catch invariant violations early.
- Around line 805-825: The gen2 dirty-card set is never cleared during minor GC
so it grows monotonically; fix it by clearing the remembered-set right after the
minor-GC fixup/promotion sweep and before re-dirtying promoted objects — e.g.
after promotion/fixup and before the loop that calls gen2_ref() +
mark_card_for_ptr(), call the gen2_cards clear/reset method (the same structure
manipulated by ensure_capacity_for_chunks and cleared during major GC) so that
stale cards are removed and only the promotion sweep/write-barrier re-dirties
needed slots; touch the code around gen2_ref, gen2_cards, mark_card_for_ptr, and
finalize_inactive_space to insert the clear at that point.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: b213ab87-7fce-41ba-9ba4-3a5e5cdb1f34
📒 Files selected for processing (28)
baml_language/crates/baml_compiler2_emit/src/lib.rsbaml_language/crates/baml_tests/tests/gc.rsbaml_language/crates/bex_engine/src/conversion.rsbaml_language/crates/bex_engine/src/heap_guard.rsbaml_language/crates/bex_engine/src/lib.rsbaml_language/crates/bex_engine/tests/early_yield.rsbaml_language/crates/bex_engine/tests/gc.rsbaml_language/crates/bex_engine/tests/heap_guard.rsbaml_language/crates/bex_engine/tests/registry.rsbaml_language/crates/bex_heap/src/accessor.rsbaml_language/crates/bex_heap/src/card_table.rsbaml_language/crates/bex_heap/src/chunked_vec.rsbaml_language/crates/bex_heap/src/gc.rsbaml_language/crates/bex_heap/src/heap.rsbaml_language/crates/bex_heap/src/heap_debugger/real.rsbaml_language/crates/bex_heap/src/heap_debugger/stub.rsbaml_language/crates/bex_heap/src/lib.rsbaml_language/crates/bex_heap/src/tlab.rsbaml_language/crates/bex_heap/tests/generational.rsbaml_language/crates/bex_vm/src/package_baml/root.rsbaml_language/crates/bex_vm/src/vm.rsbaml_language/crates/bex_vm/src/watch.rsbaml_language/crates/bex_vm/tests/early_yield.rsbaml_language/crates/bex_vm/tests/forward_roots.rsbaml_language/crates/bex_vm_types/src/lib.rsbaml_language/crates/bex_vm_types/src/roots.rsbaml_language/crates/bex_vm_types/src/types.rsbaml_language/crates/tools_onionskin/src/compiler.rs
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
baml_language/crates/bex_heap/src/card_table.rs (1)
45-55: Clarify the "Indexing" doc paragraph.The first sentence ("card
icovers objects at positions[i * CARD_SIZE, (i + 1) * CARD_SIZE)within that chunk") is only correct whilei < CARDS_PER_CHUNK. For anyi >= CARDS_PER_CHUNK, positionsi * CARD_SIZEare past the end of chunk 0 and don't live "within that chunk" at all — the precise formula a few lines below is what actually holds. Minor readability nit; consider dropping the first paragraph or rewording so it matches thechunk_idx * CARDS_PER_CHUNK + offset_in_chunk / CARD_SIZElayout.✏️ Suggested rewording
/// # Indexing /// -/// Card indices are laid out linearly: for a space with `C` chunks, card `i` -/// covers objects at positions `[i * CARD_SIZE, (i + 1) * CARD_SIZE)` within -/// that chunk. -/// -/// More precisely, for chunk index `chunk_idx` and within-chunk offset -/// `offset_in_chunk`, the card index is: +/// Card indices are laid out linearly across chunks: card `i` belongs to +/// chunk `i / CARDS_PER_CHUNK` and covers the within-chunk object range +/// `[(i % CARDS_PER_CHUNK) * CARD_SIZE, ((i % CARDS_PER_CHUNK) + 1) * CARD_SIZE)`. +/// +/// Equivalently, for chunk index `chunk_idx` and within-chunk offset +/// `offset_in_chunk`: /// ```text /// card_index = chunk_idx * CARDS_PER_CHUNK + offset_in_chunk / CARD_SIZE /// ```🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@baml_language/crates/bex_heap/src/card_table.rs` around lines 45 - 55, The first sentence of the "Indexing" doc is misleading for indices >= CARDS_PER_CHUNK; update the paragraph so it no longer states "card i covers objects at positions [i * CARD_SIZE, (i + 1) * CARD_SIZE) within that chunk" unconditionally and instead explain indexing in terms of chunk_idx and offset_in_chunk using the precise formula card_index = chunk_idx * CARDS_PER_CHUNK + offset_in_chunk / CARD_SIZE; reference the constants CARD_SIZE and CARDS_PER_CHUNK and the variables chunk_idx, offset_in_chunk and card_index in the documentation for clarity (either remove the incorrect sentence or reword it to state that that range only holds for i < CARDS_PER_CHUNK).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@baml_language/crates/bex_engine/src/heap_guard.rs`:
- Line 20: The CONST MAX_PERMITS currently set to u32::MAX can exceed Tokio's
semaphore limit on 32-bit targets and cause HeapPermitManager::new() to panic;
change the constant to a safe, target-aware value (for example use Tokio's
documented limit by computing (usize::MAX >> 3) and capping it to u32::MAX or
use tokio::sync::Semaphore::MAX_PERMITS converted to u32), update the semaphore
construction in HeapPermitManager::new() to use that capped value, and ensure
acquire_many calls use the same capped limit; if you prefer not to change
behavior, alternatively add crate metadata/docs stating the crate is
64-bit-only.
---
Nitpick comments:
In `@baml_language/crates/bex_heap/src/card_table.rs`:
- Around line 45-55: The first sentence of the "Indexing" doc is misleading for
indices >= CARDS_PER_CHUNK; update the paragraph so it no longer states "card i
covers objects at positions [i * CARD_SIZE, (i + 1) * CARD_SIZE) within that
chunk" unconditionally and instead explain indexing in terms of chunk_idx and
offset_in_chunk using the precise formula card_index = chunk_idx *
CARDS_PER_CHUNK + offset_in_chunk / CARD_SIZE; reference the constants CARD_SIZE
and CARDS_PER_CHUNK and the variables chunk_idx, offset_in_chunk and card_index
in the documentation for clarity (either remove the incorrect sentence or reword
it to state that that range only holds for i < CARDS_PER_CHUNK).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 1a5cf139-7759-4eda-b743-be1faf33e1a9
📒 Files selected for processing (9)
baml_language/crates/bex_engine/src/heap_guard.rsbaml_language/crates/bex_engine/tests/registry.rsbaml_language/crates/bex_heap/src/card_table.rsbaml_language/crates/bex_heap/src/chunked_vec.rsbaml_language/crates/bex_heap/src/heap.rsbaml_language/crates/bex_heap/src/tlab.rsbaml_language/crates/bex_heap/tests/generational.rsbaml_language/crates/bex_vm_types/src/lib.rsbaml_language/crates/tools_onionskin/src/compiler.rs
🚧 Files skipped from review as they are similar to previous changes (4)
- baml_language/crates/bex_heap/src/chunked_vec.rs
- baml_language/crates/bex_engine/tests/registry.rs
- baml_language/crates/bex_heap/src/tlab.rs
- baml_language/crates/bex_heap/src/heap.rs
Eventually we want to use dynamically sized types for heap objects, but for now we can reduce memory waste by just indirecting and having the large heap `Object` variants contain a box pointer.
Changes to Gen2, Gen1, Gen0, and inactive (scratch space) semi-spaces. The core algorithm has not really been updated yet, nor have the synchronization bugs been fixed.
Replaces old epoch system with the heap permit system. The heap permit system enforces that we either have a single exclusive heap access (`HeapGuard`) or any number of tracked active heap permits (`ActiveHeapPermit`). - Typestates `ActiveHeapPermit` and `InactiveHeapPermit` ensure we can only run VMs and other heap-root-havers when holding a permit. - `HeapGuard` guarantees that we have exclusive access to the heap (and provides heap roots). - Additional coordination takes place with `BexEngine::checking_gc` which ensures only one VM checks if we need a GC at a time so we don't have multiple GC-related park requests active or in the permit queue at the same time. - Updated the GC and function calling to use the new heap guard system.
This solves the potential issue of long-running non-yielding VM executions blocking a GC or other request to park. The reason that would be bad is that a) no protection against heap growth, and b) it blocks all other threads since they will be waiting for the GC to complete (and the GC is waiting on the long-running thread). The checks only happen in control flow instructions. Since even the simplest infinite loop will always contain control flow instructions, this is sufficient. By limiting the checks to control flow instructions, we can avoid overhead in most instructions.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
We now have a generational garbage collector. The old GC assumed single-threadedness, which is no longer the case, so failed to update other VM threads in some cases. This should no longer be an issue.
Summary by CodeRabbit
New Features
Bug Fixes / Stability
Tests