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: Shopify#552

[Bug #20209]

Author: XrXr

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")

yjit/src/codegen.rs

@@ -7258,6 +7258,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();
@@ -8038,6 +8049,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() {

yjit/src/cruby_bindings.inc.rs

@@ -975,6 +975,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;

yjit/src/stats.rs

@@ -319,6 +319,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,
@@ -405,6 +406,7 @@ make_counters! {
     invokesuper_no_me,
     invokesuper_not_iseq_or_cfunc,
     invokesuper_refinement,
+    invokesuper_singleton_class,
 
     invokeblock_megamorphic,
     invokeblock_none,