Fix false-positive memory leak using Valgrind in YJIT (#12057)

When we run with RUBY_FREE_AT_EXIT, there's a false-positive memory leak
reported in YJIT because the METHOD_CODEGEN_TABLE is never freed. This
commit adds rb_yjit_free_at_exit that is called at shutdown when
RUBY_FREE_AT_EXIT is set.

Reported memory leak:

    ==699816== 1,104 bytes in 1 blocks are possibly lost in loss record 1 of 1
    ==699816==    at 0x484680F: malloc (vg_replace_malloc.c:446)
    ==699816==    by 0x155B3E: UnknownInlinedFun (unix.rs:14)
    ==699816==    by 0x155B3E: UnknownInlinedFun (stats.rs:36)
    ==699816==    by 0x155B3E: UnknownInlinedFun (stats.rs:27)
    ==699816==    by 0x155B3E: alloc (alloc.rs:98)
    ==699816==    by 0x155B3E: alloc_impl (alloc.rs:181)
    ==699816==    by 0x155B3E: allocate (alloc.rs:241)
    ==699816==    by 0x155B3E: do_alloc<alloc::alloc::Global> (alloc.rs:15)
    ==699816==    by 0x155B3E: new_uninitialized<alloc::alloc::Global> (mod.rs:1750)
    ==699816==    by 0x155B3E: fallible_with_capacity<alloc::alloc::Global> (mod.rs:1788)
    ==699816==    by 0x155B3E: prepare_resize<alloc::alloc::Global> (mod.rs:2864)
    ==699816==    by 0x155B3E: resize_inner<alloc::alloc::Global> (mod.rs:3060)
    ==699816==    by 0x155B3E: reserve_rehash_inner<alloc::alloc::Global> (mod.rs:2950)
    ==699816==    by 0x155B3E: hashbrown::raw::RawTable<T,A>::reserve_rehash (mod.rs:1231)
    ==699816==    by 0x5BC39F: UnknownInlinedFun (mod.rs:1179)
    ==699816==    by 0x5BC39F: find_or_find_insert_slot<(usize, fn(&mut yjit::codegen::JITState, &mut yjit::backend::ir::Assembler, *const yjit::cruby::autogened::rb_callinfo, *const yjit::cruby::autogened::rb_callable_method_entry_struct, core::option::Option<yjit::codegen::BlockHandler>, i32, core::option::Option<yjit::cruby::VALUE>) -> bool), alloc::alloc::Global, hashbrown::map::equivalent_key::{closure_env#0}<usize, usize, fn(&mut yjit::codegen::JITState, &mut yjit::backend::ir::Assembler, *const yjit::cruby::autogened::rb_callinfo, *const yjit::cruby::autogened::rb_callable_method_entry_struct, core::option::Option<yjit::codegen::BlockHandler>, i32, core::option::Option<yjit::cruby::VALUE>) -> bool>, hashbrown::map::make_hasher::{closure_env#0}<usize, fn(&mut yjit::codegen::JITState, &mut yjit::backend::ir::Assembler, *const yjit::cruby::autogened::rb_callinfo, *const yjit::cruby::autogened::rb_callable_method_entry_struct, core::option::Option<yjit::codegen::BlockHandler>, i32, core::option::Option<yjit::cruby::VALUE>) -> bool, std:#️⃣:random::RandomState>> (mod.rs:1413)
    ==699816==    by 0x5BC39F: hashbrown::map::HashMap<K,V,S,A>::insert (map.rs:1754)
    ==699816==    by 0x57C5C6: insert<usize, fn(&mut yjit::codegen::JITState, &mut yjit::backend::ir::Assembler, *const yjit::cruby::autogened::rb_callinfo, *const yjit::cruby::autogened::rb_callable_method_entry_struct, core::option::Option<yjit::codegen::BlockHandler>, i32, core::option::Option<yjit::cruby::VALUE>) -> bool, std:#️⃣:random::RandomState> (map.rs:1104)
    ==699816==    by 0x57C5C6: yjit::codegen::reg_method_codegen (codegen.rs:10521)
    ==699816==    by 0x57C295: yjit::codegen::yjit_reg_method_codegen_fns (codegen.rs:10464)
    ==699816==    by 0x5C6B07: rb_yjit_init (yjit.rs:40)
    ==699816==    by 0x393723: ruby_opt_init (ruby.c:1820)
    ==699816==    by 0x393723: ruby_opt_init (ruby.c:1767)
    ==699816==    by 0x3957D4: prism_script (ruby.c:2215)
    ==699816==    by 0x3957D4: process_options (ruby.c:2538)
    ==699816==    by 0x396065: ruby_process_options (ruby.c:3166)
    ==699816==    by 0x236E56: ruby_options (eval.c:117)
    ==699816==    by 0x15BAED: rb_main (main.c:43)
    ==699816==    by 0x15BAED: main (main.c:62)

After this patch, there are no more memory leaks reported when running
RUBY_FREE_AT_EXIT with Valgrind on an empty Ruby script:

    $ RUBY_FREE_AT_EXIT=1 valgrind --leak-check=full ruby -e ""
    ...
    ==700357== HEAP SUMMARY:
    ==700357==     in use at exit: 0 bytes in 0 blocks
    ==700357==   total heap usage: 36,559 allocs, 36,559 frees, 6,064,783 bytes allocated
    ==700357==
    ==700357== All heap blocks were freed -- no leaks are possible

vm.c

@@ -3161,6 +3161,12 @@ ruby_vm_destruct(rb_vm_t *vm)
         /* after freeing objspace, you *can't* use ruby_xfree() */
         ruby_mimfree(vm);
         ruby_current_vm_ptr = NULL;
+
+#if USE_YJIT
+        if (rb_free_at_exit) {
+            rb_yjit_free_at_exit();
+        }
+#endif
     }
     RUBY_FREE_LEAVE("vm");
     return 0;

yjit.h

@@ -37,6 +37,7 @@ void rb_yjit_collect_binding_alloc(void);
 void rb_yjit_collect_binding_set(void);
 void rb_yjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception);
 void rb_yjit_init(bool yjit_enabled);
+void rb_yjit_free_at_exit();
 void rb_yjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop);
 void rb_yjit_constant_state_changed(ID id);
 void rb_yjit_iseq_mark(void *payload);

yjit/src/codegen.rs

@@ -10521,6 +10521,10 @@ fn reg_method_codegen(klass: VALUE, mid_str: &str, gen_fn: MethodGenFn) {
     unsafe { METHOD_CODEGEN_TABLE.as_mut().unwrap().insert(method_serial, gen_fn); }
 }
 
+pub fn yjit_shutdown_free_codegen_table() {
+    unsafe { METHOD_CODEGEN_TABLE = None; };
+}
+
 /// Global state needed for code generation
 pub struct CodegenGlobals {
     /// Flat vector of bits to store compressed context data

yjit/src/yjit.rs

@@ -78,6 +78,11 @@ fn yjit_init() {
     }
 }
 
+#[no_mangle]
+pub extern "C" fn rb_yjit_free_at_exit() {
+    yjit_shutdown_free_codegen_table();
+}
+
 /// At the moment, we abort in all cases we panic.
 /// To aid with getting diagnostics in the wild without requiring
 /// people to set RUST_BACKTRACE=1, register a panic hook that crash using rb_bug().