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

2 files changed, 9 insertions, 11 deletions

diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
index ed6feb3174..0320fdd829 100644
--- a/yjit/src/asm/mod.rs
+++ b/yjit/src/asm/mod.rs
@@ -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);
     }
diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index b695f8da96..66e333f867 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -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
         };