From 7442cb461b32de2eec3b37f52d80752d30627de0 Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Mon, 7 Nov 2022 07:48:26 -0800
Subject: [PATCH] YJIT: Free pages after ObjectSpace API usages (#6676)

---
 yjit/src/asm/mod.rs | 28 +++++++++++++++-------------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
index 3d7de7cd79..497f7687ed 100644
--- a/yjit/src/asm/mod.rs
+++ b/yjit/src/asm/mod.rs
@@ -563,23 +563,14 @@ impl CodeBlock {
                 pages_in_use[*page] = true;
             }
         });
-        // Outlined code generated by CodegenGlobals::init() should also be kept.
-        for page in CodegenGlobals::get_ocb_pages() {
-            pages_in_use[*page] = true;
-        }
-
-        // Let VirtuamMem free the pages
-        let mut freed_pages: Vec<usize> = pages_in_use.iter().enumerate()
-            .filter(|&(_, &in_use)| !in_use).map(|(page, _)| page).collect();
-        self.free_pages(&freed_pages);
         // Avoid accumulating freed pages for future code GC
         for_each_off_stack_iseq_payload(|iseq_payload: &mut IseqPayload| {
             iseq_payload.pages.clear();
         });
-
-        // Append virtual pages in case RubyVM::YJIT.code_gc is manually triggered.
-        let mut virtual_pages: Vec<usize> = (self.num_mapped_pages()..self.num_virtual_pages()).collect();
-        freed_pages.append(&mut virtual_pages);
+        // Outlined code generated by CodegenGlobals::init() should also be kept.
+        for page in CodegenGlobals::get_ocb_pages() {
+            pages_in_use[*page] = true;
+        }
 
         // Invalidate everything to have more compact code after code GC.
         // This currently patches every ISEQ, which works, but in the future,
@@ -591,6 +582,17 @@ impl CodeBlock {
         // can be safely reset to pass the frozen bytes check on invalidation.
         CodegenGlobals::set_inline_frozen_bytes(0);
 
+        // Let VirtuamMem free the pages
+        let mut freed_pages: Vec<usize> = pages_in_use.iter().enumerate()
+            .filter(|&(_, &in_use)| !in_use).map(|(page, _)| page).collect();
+        // ObjectSpace API may trigger Ruby's GC, which marks gc_offsets in JIT code.
+        // So this should be called after for_each_*_iseq_payload and rb_yjit_tracing_invalidate_all.
+        self.free_pages(&freed_pages);
+
+        // Append virtual pages in case RubyVM::YJIT.code_gc is manually triggered.
+        let mut virtual_pages: Vec<usize> = (self.num_mapped_pages()..self.num_virtual_pages()).collect();
+        freed_pages.append(&mut virtual_pages);
+
         if let Some(&first_page) = freed_pages.first() {
             let mut cb = CodegenGlobals::get_inline_cb();
             cb.write_pos = cb.get_page_pos(first_page);