ZJIT: Restore dropped_bytes after temporary OOM (ruby#15069)

Author: k0kubun

zjit/src/asm/mod.rs

@@ -208,6 +208,14 @@ impl CodeBlock {
         self.dropped_bytes
     }
 
+    /// Set dropped_bytes to false if the current zjit_alloc_bytes() + code_region_size
+    /// + page_size is below --zjit-mem-size.
+    pub fn update_dropped_bytes(&mut self) {
+        if self.mem_block.borrow().can_allocate() {
+            self.dropped_bytes = false;
+        }
+    }
+
     /// Allocate a new label with a given name
     pub fn new_label(&mut self, name: String) -> Label {
         assert!(!name.contains(' '), "use underscores in label names, not spaces");

zjit/src/backend/lir.rs

@@ -1750,7 +1750,15 @@ impl Assembler
         #[cfg(feature = "disasm")]
         let start_addr = cb.get_write_ptr();
         let alloc_regs = Self::get_alloc_regs();
-        let ret = self.compile_with_regs(cb, alloc_regs);
+        let had_dropped_bytes = cb.has_dropped_bytes();
+        let ret = self.compile_with_regs(cb, alloc_regs).inspect_err(|err| {
+            // If we use too much memory to compile the Assembler, it would set cb.dropped_bytes = true.
+            // To avoid failing future compilation by cb.has_dropped_bytes(), attempt to reset dropped_bytes with
+            // the current zjit_alloc_bytes() which may be decreased after self is dropped in compile_with_regs().
+            if *err == CompileError::OutOfMemory && !had_dropped_bytes {
+                cb.update_dropped_bytes();
+            }
+        });
 
         #[cfg(feature = "disasm")]
         if get_option!(dump_disasm) {

zjit/src/codegen.rs

@@ -220,6 +220,11 @@ fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr, function: Option<&Function>) -> R
 
 /// Compile an ISEQ into machine code
 fn gen_iseq_body(cb: &mut CodeBlock, iseq: IseqPtr, function: Option<&Function>, payload: &mut IseqPayload) -> Result<IseqCodePtrs, CompileError> {
+    // If we ran out of code region, we shouldn't attempt to generate new code.
+    if cb.has_dropped_bytes() {
+        return Err(CompileError::OutOfMemory);
+    }
+
     // Convert ISEQ into optimized High-level IR if not given
     let function = match function {
         Some(function) => function,

zjit/src/virtualmem.rs

@@ -258,6 +258,13 @@ impl<A: Allocator> VirtualMemory<A> {
         Ok(())
     }
 
+    /// Return true if write_byte() can allocate a new page
+    pub fn can_allocate(&self) -> bool {
+        let memory_usage_bytes = self.mapped_region_bytes + zjit_alloc_bytes();
+        let memory_limit_bytes = self.memory_limit_bytes.unwrap_or(self.region_size_bytes);
+        memory_usage_bytes + self.page_size_bytes < memory_limit_bytes
+    }
+
     /// Make all the code in the region executable. Call this at the end of a write session.
     /// See [Self] for usual usage flow.
     pub fn mark_all_executable(&mut self) {