Fix concurrent escape bug. Comment about over spilling

Author: XrXr

vm.c

@@ -1146,10 +1146,14 @@ vm_make_env_each(const rb_execution_context_t * const ec, rb_control_frame_t *co
     if (VM_FRAME_RUBYFRAME_P(cfp) &&
         !rbimpl_atomic_load(&ISEQ_BODY(iseq)->jit_ep_escape_recorded, RBIMPL_ATOMIC_RELAXED)) {
         if (rb_yjit_enabled_p) rb_yjit_invalidate_ep_is_bp(iseq);
-        if (rb_zjit_enabled_p) {
-            rb_zjit_invalidate_no_ep_escape(iseq);
-            rb_zjit_spill_frame(cfp);
-        }
+        if (rb_zjit_enabled_p) rb_zjit_invalidate_no_ep_escape(iseq);
+    }
+
+    // Spill ZJIT frame if there is one. This is a separate concern from
+    // invalidation and needs to happen in all execution threads that reach here.
+    if (VM_FRAME_RUBYFRAME_P(cfp)) {
+        // We only need the locals here but we reconstruct the stack as well for simplicity.
+        rb_zjit_spill_frame(cfp);
     }
 
     /*

zjit/src/codegen_tests.rs

@@ -8239,6 +8239,44 @@ fn test_no_ep_escape_invalidation_at_max_versions() {
     assert_snapshot!(result, @r#""expected""#);
 }
 
+#[test]
+fn test_concurrent_ep_escape() {
+    // Test to see if concurrent environment escapes
+    // sharing one piece of JIT code produce one consistent result.
+    rb_zjit_prepare_options();
+    set_call_threshold(2);
+    set_inline_threshold(0);
+    let result = inspect(r#"
+        def spill_escaper(gate)
+          a = 1
+          b = 2
+          spill_captor(gate) { }  # empty block, no eager spill
+        end
+
+        # Returning `blk` materializes the Proc, which escapes spill_escaper's EP.
+        def spill_captor(gate, &blk)
+          gate.receive ? blk : nil
+        end
+
+        ready = Ractor::Port.new
+        results = Ractor::Port.new
+        ractors = 4.times.map do
+          Ractor.new(ready, results) do |ready, results|
+            port = Ractor.current.default_port
+            2.times { port << false; spill_escaper(port) }  # compile spill_escaper, no escape
+            ready << :parked                                # parks until the main Ractor releases
+            bnd = spill_escaper(port).binding
+            results << [bnd.local_variable_get(:a), bnd.local_variable_get(:b)]
+          end
+        end
+
+        ractors.size.times { ready.receive }
+        ractors.each { |r| r << true }  # have all ractor do EP escapes
+        ractors.size.times.map { results.receive }.uniq
+    "#);
+    assert_snapshot!(result, @r#"[[1, 2]]"#);
+}
+
 #[test]
 fn test_float_arithmetic() {
     set_call_threshold(1);