YJIT: A64: Remove assert that trips when OOM at page boundary

With a well-timed OOM around a page switch in the backend, it can return
RetryOnNextPage twice and crash due to the assert. (More places can
signal OOM now since VirtualMem tracks Rust malloc heap size for
--yjit-mem-size.)

Return error in these cases instead of crashing.

Fixes: https://github.com/Shopify/ruby/issues/566
This commit is contained in:
Alan Wu 2025-01-29 14:59:26 -05:00
parent de45755de8
commit 5a7089fc03
Notes: git 2025-01-30 00:09:55 +00:00
2 changed files with 9 additions and 11 deletions

View File

@ -145,6 +145,7 @@ impl CodeBlock {
/// Move the CodeBlock to the next page. If it's on the furthest page,
/// move the other CodeBlock to the next page as well.
#[must_use]
pub fn next_page<F: Fn(&mut CodeBlock, CodePtr)>(&mut self, base_ptr: CodePtr, jmp_ptr: F) -> bool {
let old_write_ptr = self.get_write_ptr();
self.set_write_ptr(base_ptr);
@ -823,7 +824,7 @@ mod tests
assert_eq!(cb.code_size(), 4);
// Moving to the next page should not increase code_size
cb.next_page(cb.get_write_ptr(), |_, _| {});
assert!(cb.next_page(cb.get_write_ptr(), |_, _| {}));
assert_eq!(cb.code_size(), 4);
// Write 4 bytes in the second page
@ -836,7 +837,7 @@ mod tests
cb.write_bytes(&[1, 1, 1, 1]);
// Moving from an old page to the next page should not increase code_size
cb.next_page(cb.get_write_ptr(), |_, _| {});
assert!(cb.next_page(cb.get_write_ptr(), |_, _| {}));
cb.set_pos(old_write_pos);
assert_eq!(cb.code_size(), 8);
}

View File

@ -1341,16 +1341,13 @@ impl Assembler
Err(EmitError::RetryOnNextPage) => {
// we want to lower jumps to labels to b.cond instructions, which have a 1 MiB
// range limit. We can easily exceed the limit in case the jump straddles two pages.
// In this case, we retry with a fresh page.
// In this case, we retry with a fresh page once.
cb.set_label_state(starting_label_state);
cb.next_page(start_ptr, emit_jmp_ptr_with_invalidation);
let result = asm.arm64_emit(cb, &mut ocb);
assert_ne!(
Err(EmitError::RetryOnNextPage),
result,
"should not fail when writing to a fresh code page"
);
result
if cb.next_page(start_ptr, emit_jmp_ptr_with_invalidation) {
asm.arm64_emit(cb, &mut ocb)
} else {
Err(EmitError::OutOfMemory)
}
}
result => result
};