Backport 3.3: YJIT memory leak fix with additional CI fixes (#9841)

merge revision(s) 2cc7a56e,b0711b1,db5d9429: [Backport #20209]

	YJIT: Avoid leaks by skipping objects with a singleton class

	For receiver with a singleton class, there are multiple vectors YJIT can
	end up retaining the object. There is a path in jit_guard_known_klass()
	that bakes the receiver into the code, and the object could also be kept
	alive indirectly through a path starting at the CME object baked into
	the code.

	To avoid these leaks, avoid compiling calls on objects with a singleton
	class.

	See: https://github.com/Shopify/ruby/issues/552

	[Bug #20209]
	---
	 yjit/bindgen/src/main.rs       |  1 +
	 yjit/src/codegen.rs            | 17 +++++++++++++++++
	 yjit/src/cruby_bindings.inc.rs |  1 +
	 yjit/src/stats.rs              |  2 ++
	 4 files changed, 21 insertions(+)

	YJIT: Fix tailcall and JIT entry eating up FINISH frames (#9729)

	Suppose YJIT runs a rb_vm_opt_send_without_block()
	fallback and the control frame stack looks like:

	```
	will_tailcall_bar [FINISH]
	caller_that_used_fallback
	```

	will_tailcall_bar() runs in the interpreter and sets up a tailcall.
	Right before JIT_EXEC() in the `send` instruction, the stack will look like:

	```
	bar [FINISH]
	caller_that_used_fallback
	```

	Previously, JIT_EXEC() ran bar() in JIT code, which caused the `FINISH`
	flag to return to the interpreter instead of to the JIT code running
	caller_that_used_fallback(), causing code to run twice and probably
	crash. Recent flaky failures on CI about "each stub expects a particular
	iseq" are probably due to leaving methods twice in
	`test_optimizations.rb`.

	Only run JIT code from the interpreter if a new frame is pushed.
	---
	 test/ruby/test_optimization.rb | 11 +++++++++++
	 vm_exec.h                      |  3 ++-
	 2 files changed, 13 insertions(+), 1 deletion(-)

	YJIT: No need to RESTORE_REG now that we reject tailcalls

	Thanks to Kokubun for noticing.

	Follow-up: b0711b1cf152afad0a480ee2f9bedd142a0d24ac
	---
	 vm_exec.h | 1 -
	 1 file changed, 1 deletion(-)

6 files changed, 34 insertions, 2 deletions

diff --git a/test/ruby/test_optimization.rb b/test/ruby/test_optimization.rb
index 8d669e502c..70b6bde6ed 100644
--- a/test/ruby/test_optimization.rb
+++ b/test/ruby/test_optimization.rb
@@ -451,6 +451,17 @@ class TestRubyOptimization < Test::Unit::TestCase
     assert_equal(3, one_plus_two)
   end
 
+  def test_tailcall_and_post_arg
+    tailcall(<<~RUBY)
+      def ret_const = :ok
+
+      def post_arg(_a = 1, _b) = ret_const
+    RUBY
+
+    # YJIT probably uses a fallback on the call to post_arg
+    assert_equal(:ok, post_arg(0))
+  end
+
   def test_tailcall_interrupted_by_sigint
     bug12576 = 'ruby-core:76327'
     script = "#{<<-"begin;"}\n#{<<~'end;'}"
diff --git a/vm_exec.h b/vm_exec.h
index 152410a6a7..11b89c30fc 100644
--- a/vm_exec.h
+++ b/vm_exec.h
@@ -176,9 +176,9 @@ default:                        \
 // Run the JIT from the interpreter
 #define JIT_EXEC(ec, val) do { \
     rb_jit_func_t func; \
-    if (val == Qundef && (func = jit_compile(ec))) { \
+    /* don't run tailcalls since that breaks FINISH */ \
+    if (val == Qundef && GET_CFP() != ec->cfp && (func = jit_compile(ec))) { \
         val = func(ec, ec->cfp); \
-        RESTORE_REGS(); /* fix cfp for tailcall */ \
         if (ec->tag->state) THROW_EXCEPTION(val); \
     } \
 } while (0)
diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 4308a4cb40..b4b4d17f5d 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -447,6 +447,7 @@ fn main() {
         .allowlist_function("rb_obj_is_proc")
         .allowlist_function("rb_vm_base_ptr")
         .allowlist_function("rb_ec_stack_check")
+        .allowlist_function("rb_vm_top_self")
 
         // We define VALUE manually, don't import it
         .blocklist_type("VALUE")
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index f360c33bcb..75986dedb1 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -7132,6 +7132,17 @@ fn gen_send_general(
     assert_eq!(RUBY_T_CLASS, comptime_recv_klass.builtin_type(),
         "objects visible to ruby code should have a T_CLASS in their klass field");
 
+    // Don't compile calls through singleton classes to avoid retaining the receiver.
+    // Make an exception for class methods since classes tend to be retained anyways.
+    // Also compile calls on top_self to help tests.
+    if VALUE(0) != unsafe { FL_TEST(comptime_recv_klass, VALUE(RUBY_FL_SINGLETON as usize)) }
+        && comptime_recv != unsafe { rb_vm_top_self() }
+        && !unsafe { RB_TYPE_P(comptime_recv, RUBY_T_CLASS) }
+        && !unsafe { RB_TYPE_P(comptime_recv, RUBY_T_MODULE) } {
+        gen_counter_incr(asm, Counter::send_singleton_class);
+        return None;
+    }
+
     // Points to the receiver operand on the stack
     let recv = asm.stack_opnd(recv_idx);
     let recv_opnd: YARVOpnd = recv.into();
@@ -7887,6 +7898,12 @@ fn gen_invokesuper_specialized(
         return None;
     }
 
+    // Don't compile `super` on objects with singleton class to avoid retaining the receiver.
+    if VALUE(0) != unsafe { FL_TEST(comptime_recv.class_of(), VALUE(RUBY_FL_SINGLETON as usize)) } {
+        gen_counter_incr(asm, Counter::invokesuper_singleton_class);
+        return None;
+    }
+
     // Do method lookup
     let cme = unsafe { rb_callable_method_entry(comptime_superclass, mid) };
     if cme.is_null() {
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 4b49ab3ed0..394c6d013c 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -956,6 +956,7 @@ extern "C" {
         n: ::std::os::raw::c_long,
         elts: *const VALUE,
     ) -> VALUE;
+    pub fn rb_vm_top_self() -> VALUE;
     pub static mut rb_vm_insns_count: u64;
     pub fn rb_method_entry_at(obj: VALUE, id: ID) -> *const rb_method_entry_t;
     pub fn rb_callable_method_entry(klass: VALUE, id: ID) -> *const rb_callable_method_entry_t;
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index ac247d2fa8..dcaa9af2b5 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -300,6 +300,7 @@ make_counters! {
     // Method calls that fallback to dynamic dispatch
     send_keywords,
     send_kw_splat,
+    send_singleton_class,
     send_args_splat_super,
     send_iseq_zsuper,
     send_block_arg,
@@ -380,6 +381,7 @@ make_counters! {
     invokesuper_no_me,
     invokesuper_not_iseq_or_cfunc,
     invokesuper_refinement,
+    invokesuper_singleton_class,
 
     invokeblock_megamorphic,
     invokeblock_none,