From 63d96ccbcd21653ac1e5afd96dbd72bc78900de0 Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Fri, 24 Feb 2023 13:19:42 -0800
Subject: [PATCH] Invalidate everything on GC.compact

---
 iseq.c                         |  2 +-
 lib/ruby_vm/mjit/hooks.rb      |  4 +++
 lib/ruby_vm/mjit/invariants.rb | 49 +++++++++++++++++++++-------------
 mjit.c                         | 34 +++++++++++++++++------
 mjit.h                         |  2 +-
 5 files changed, 62 insertions(+), 29 deletions(-)

diff --git a/iseq.c b/iseq.c
index 3305edde96..03a7336455 100644
--- a/iseq.c
+++ b/iseq.c
@@ -357,7 +357,7 @@ rb_iseq_mark_and_move(rb_iseq_t *iseq, bool reference_updating)
 
         if (reference_updating) {
 #if USE_MJIT
-            mjit_update_references(iseq);
+            rb_mjit_iseq_update_references(iseq);
 #endif
 #if USE_YJIT
             rb_yjit_iseq_update_references(body->yjit_payload);
diff --git a/lib/ruby_vm/mjit/hooks.rb b/lib/ruby_vm/mjit/hooks.rb
index f3943504fd..477c4e0ae8 100644
--- a/lib/ruby_vm/mjit/hooks.rb
+++ b/lib/ruby_vm/mjit/hooks.rb
@@ -26,5 +26,9 @@ module RubyVM::MJIT
     def self.on_tracing_invalidate_all(_new_iseq_events)
       Invariants.on_tracing_invalidate_all
     end
+
+    def self.on_update_references
+      Invariants.on_update_references
+    end
   end
 end
diff --git a/lib/ruby_vm/mjit/invariants.rb b/lib/ruby_vm/mjit/invariants.rb
index e68a6e5503..0e8063c0c6 100644
--- a/lib/ruby_vm/mjit/invariants.rb
+++ b/lib/ruby_vm/mjit/invariants.rb
@@ -78,26 +78,11 @@ module RubyVM::MJIT
       end
 
       def on_tracing_invalidate_all
-        # On-Stack Replacement
-        @patches.each do |address, target|
-          # TODO: assert patches don't overlap each other
-          @cb.with_write_addr(address) do
-            asm = Assembler.new
-            asm.comment('on_tracing_invalidate_all')
-            asm.jmp(target)
-            @cb.write(asm)
-          end
-        end
+        invalidate_all
+      end
 
-        # Avoid reusing past code
-        Compiler.reset_blocks
-
-        C.mjit_for_each_iseq do |iseq|
-          # Disable entering past code
-          iseq.body.jit_func = 0
-          # Compile this again if not converted to trace_* insns
-          iseq.body.total_calls = 0
-        end
+      def on_update_references
+        invalidate_all
       end
 
       # @param jit [RubyVM::MJIT::JITState]
@@ -111,6 +96,32 @@ module RubyVM::MJIT
           end
         end
       end
+
+      private
+
+      def invalidate_all
+        # On-Stack Replacement
+        @patches.each do |address, target|
+          # TODO: assert patches don't overlap each other
+          @cb.with_write_addr(address) do
+            asm = Assembler.new
+            asm.comment('on_tracing_invalidate_all')
+            asm.jmp(target)
+            @cb.write(asm)
+          end
+        end
+        @patches.clear
+
+        # Avoid reusing past code
+        Compiler.reset_blocks
+
+        C.mjit_for_each_iseq do |iseq|
+          # Disable entering past code
+          iseq.body.jit_func = 0
+          # Compile this again if not converted to trace_* insns
+          iseq.body.total_calls = 0
+        end
+      end
     end
   end
 end
diff --git a/mjit.c b/mjit.c
index de55070c58..c140a5f252 100644
--- a/mjit.c
+++ b/mjit.c
@@ -113,12 +113,6 @@ mjit_cancel_all(const char *reason)
     }
 }
 
-void
-mjit_update_references(const rb_iseq_t *iseq)
-{
-    // TODO: remove this
-}
-
 void
 mjit_free_iseq(const rb_iseq_t *iseq)
 {
@@ -348,12 +342,13 @@ mjit_cme_invalidate(void *data)
     });
 }
 
+extern int rb_workqueue_register(unsigned flags, rb_postponed_job_func_t func, void *data);
+
 void
 rb_mjit_cme_invalidate(rb_callable_method_entry_t *cme)
 {
     if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
     // Asynchronously hook the Ruby code since running Ruby in the middle of cme invalidation is dangerous.
-    extern int rb_workqueue_register(unsigned flags, rb_postponed_job_func_t func, void *data);
     rb_workqueue_register(0, mjit_cme_invalidate, (void *)cme);
 }
 
@@ -389,6 +384,28 @@ rb_mjit_tracing_invalidate_all(rb_event_flag_t new_iseq_events)
     });
 }
 
+static void
+mjit_iseq_update_references(void *data)
+{
+    if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
+    WITH_MJIT_ISOLATED({
+        rb_funcall(rb_mMJITHooks, rb_intern("on_update_references"), 0);
+    });
+}
+
+void
+rb_mjit_iseq_update_references(const rb_iseq_t *iseq)
+{
+    if (!mjit_enabled) return;
+
+    // TODO: update mjit_blocks
+
+    // Asynchronously hook the Ruby code to avoid allocation during GC.compact.
+    // Using _one because it's too slow to invalidate all for each ISEQ. Thus
+    // not giving an ISEQ pointer.
+    rb_postponed_job_register_one(0, mjit_iseq_update_references, NULL);
+}
+
 // TODO: Use this in more places
 VALUE
 rb_mjit_iseq_new(rb_iseq_t *iseq)
@@ -442,10 +459,11 @@ mjit_mark(void)
         return;
     RUBY_MARK_ENTER("mjit");
 
-    // Mark objects used by the MJIT compiler
+    // Pin object pointers used in this file
     rb_gc_mark(rb_MJITCompiler);
     rb_gc_mark(rb_cMJITIseqPtr);
     rb_gc_mark(rb_cMJITCfpPtr);
+    rb_gc_mark(rb_mMJITHooks);
 
     RUBY_MARK_LEAVE("mjit");
 }
diff --git a/mjit.h b/mjit.h
index 3f7719d51d..3b65911009 100644
--- a/mjit.h
+++ b/mjit.h
@@ -106,7 +106,7 @@ extern void mjit_cancel_all(const char *reason);
 extern bool mjit_compile(FILE *f, const rb_iseq_t *iseq, const char *funcname, int id);
 extern void mjit_init(const struct mjit_options *opts);
 extern void mjit_free_iseq(const rb_iseq_t *iseq);
-extern void mjit_update_references(const rb_iseq_t *iseq);
+extern void rb_mjit_iseq_update_references(const rb_iseq_t *iseq);
 extern void mjit_mark(void);
 extern void mjit_mark_cc_entries(const struct rb_iseq_constant_body *const body);
 extern void mjit_notify_waitpid(int exit_code);