merge revision(s) 8f98abfc46d48c84db2b1408fc8f14b240ec05fd: [Backport #21941]

	[PATCH] YJIT: Fix not reading locals from `cfp->ep` after `YJIT.enable` and exceptional entry
	MIME-Version: 1.0
	Content-Type: text/plain; charset=UTF-8
	Content-Transfer-Encoding: 8bit

	Fix for [Bug #21941].

	In case of `--yjit-disable`, YJIT only starts to record environment
	escapes after `RubyVM::YJIT.enable`. Previously we falsely assumed that
	we always have a full history all the way back to VM boot. This had YJIT
	install and run code that assume EP=BP when EP≠BP for some exceptional
	entry into the middle of a running frame, if the environment escaped
	before `YJIT.enable`.

	The fix is to reject exceptional entry with an escaped environment.
	Rename things and explain in more detail how the predicate for deciding
	to assume EP=BP works. It's quite subtle since it reasons about all
	parties in the system that push a control frame and then run JIT code.

	Note that while can_assume_on_stack_env() checks the currently running
	environment if it so happens to be the one YJIT is compiling against, it
	can return true for any ISEQ. The check isn't necessary for fixing the
	bug, and the load bearing part of this patch is the change to
	exceptional entries.

	This fix is flat on speed and space on ruby-bench headline benchmarks.

	Many thanks for the community effort to create a small test case for
	this bug.

8 files changed, 74 insertions, 16 deletions

diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index cf348b778d..962b329d50 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -1762,6 +1762,27 @@ class TestYJIT < Test::Unit::TestCase
     RUBY
   end
 
+  def test_exceptional_entry_into_env_escaped_before_yjit_enablement
+    threshold = 2
+    assert_separately(["--disable-all", "--yjit-disable", "--yjit-call-threshold=#{threshold}"], <<~RUBY)
+      def run
+        @captured_env = ->{}
+        RubyVM::YJIT.enable
+
+        i = 0
+        while i < #{threshold}
+          next_i = i + 1
+          from_break = tap { break i + 1 } # break from the block generates an exceptional entry
+          assert_equal(from_break, next_i, '[Bug #21941]')
+          i = next_i
+        end
+      end
+
+      run
+      assert_equal(#{threshold}, @captured_env.binding.local_variable_get(:i))
+    RUBY
+  end
+
   private
 
   def code_gc_helpers
diff --git a/version.h b/version.h
index f53bbdd8b3..46182a0b6d 100644
--- a/version.h
+++ b/version.h
@@ -11,7 +11,7 @@
 # define RUBY_VERSION_MINOR RUBY_API_VERSION_MINOR
 #define RUBY_VERSION_TEENY 9
 #define RUBY_RELEASE_DATE RUBY_RELEASE_YEAR_STR"-"RUBY_RELEASE_MONTH_STR"-"RUBY_RELEASE_DAY_STR
-#define RUBY_PATCHLEVEL 84
+#define RUBY_PATCHLEVEL 85
 
 #include "ruby/version.h"
 #include "ruby/internal/abi.h"
diff --git a/yjit.rb b/yjit.rb
index 556ecf25f6..ab36e36da0 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -338,6 +338,9 @@ module RubyVM::YJIT
       # Number of failed compiler invocations
       compilation_failure = stats[:compilation_failure]
 
+      # Number of refused exceptional entries with an escaped environment
+      exceptional_entry_escaped_env = stats[:exceptional_entry_escaped_env]
+
       code_region_overhead = stats[:code_region_size] - (stats[:inline_code_size] + stats[:outlined_code_size])
 
       out.puts "num_send:              " + format_number(13, stats[:num_send])
@@ -374,6 +377,7 @@ module RubyVM::YJIT
       out.puts "bindings_allocations:  " + format_number(13, stats[:binding_allocations])
       out.puts "bindings_set:          " + format_number(13, stats[:binding_set])
       out.puts "compilation_failure:   " + format_number(13, compilation_failure) if compilation_failure != 0
+      out.puts "exceptional_entry_escaped_env:" + format_number(6, exceptional_entry_escaped_env) if exceptional_entry_escaped_env != 0
       out.puts "live_iseq_count:       " + format_number(13, stats[:live_iseq_count])
       out.puts "iseq_alloc_count:      " + format_number(13, stats[:iseq_alloc_count])
       out.puts "compiled_iseq_entry:   " + format_number(13, stats[:compiled_iseq_entry])
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 437e5d1505..8c284469b3 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -234,20 +234,36 @@ impl<'a> JITState<'a> {
         result
     }
 
-    /// Return true if the current ISEQ could escape an environment.
+    /// Return true if the JIT code can use [`Self::assume_no_ep_escape`]
+    /// and will run with an on-stack (`!VM_ENV_ESCAPED_P`) environment.
     ///
-    /// As of vm_push_frame(), EP is always equal to BP. However, after pushing
-    /// a frame, some ISEQ setups call vm_bind_update_env(), which redirects EP.
-    /// Also, some method calls escape the environment to the heap.
-    fn escapes_ep(&self) -> bool {
+    /// ## Reasoning about ISEQs that are not currently running
+    ///
+    /// As of vm_push_frame() and its JIT code equivalent, EP is always equal to BP (the
+    /// environment is on-stack and has not escaped). We can usually assume this is the starting
+    /// condition upon entry into JIT code. However, after pushing a frame and before entry into
+    /// JIT code, some ISEQ setups call vm_bind_update_env(), which redirects EP.
+    ///
+    /// ## After making the assumption
+    ///
+    /// After JIT code entry, many ruby operations can have the environment escape to heap. These
+    /// are handled by [`crate::invariants`].
+    ///
+    /// Exceptional entry through jit_exec_exception() is an extreme case of the environment state
+    /// changing between vm_push_frame() and entry into JIT code. The frame could have been pushed
+    /// before YJIT is enabled. The exception entry point refuses entry with an escaped environment.
+    fn can_assume_on_stack_env(&self) -> bool {
         match unsafe { get_iseq_body_type(self.iseq) } {
             // <main> frame is always associated to TOPLEVEL_BINDING.
             ISEQ_TYPE_MAIN |
             // Kernel#eval uses a heap EP when a Binding argument is not nil.
-            ISEQ_TYPE_EVAL => true,
-            // If this ISEQ has previously escaped EP, give up the optimization.
-            _ if iseq_escapes_ep(self.iseq) => true,
-            _ => false,
+            ISEQ_TYPE_EVAL => false,
+            // Check the running environment if compiling for it
+            _ if unsafe { self.iseq == get_cfp_iseq(self.get_cfp()) && cfp_env_has_escaped(self.get_cfp()) } => false,
+            // If we've seen this ISEQ run with an escaped environment, give up the optimization
+            // to avoid excessive invalidations (even though it may be fine for soundness).
+            _ if seen_escaped_env(self.iseq) => false,
+            _ => true,
         }
     }
 
@@ -376,8 +392,8 @@ impl<'a> JITState<'a> {
         if jit_ensure_block_entry_exit(self, asm).is_none() {
             return false; // out of space, give up
         }
-        if self.escapes_ep() {
-            return false; // EP has been escaped in this ISEQ. disable the optimization to avoid an invalidation loop.
+        if !self.can_assume_on_stack_env() {
+            return false; // Unsound or unprofitable to make the assumption
         }
         self.no_ep_escape = true;
         true
@@ -2448,7 +2464,7 @@ fn gen_getlocal_generic(
     level: u32,
 ) -> Option<CodegenStatus> {
     // Split the block if we need to invalidate this instruction when EP escapes
-    if level == 0 && !jit.escapes_ep() && !jit.at_compile_target() {
+    if level == 0 && jit.can_assume_on_stack_env() && !jit.at_compile_target() {
         return jit.defer_compilation(asm);
     }
 
@@ -2549,7 +2565,7 @@ fn gen_setlocal_generic(
     }
 
     // Split the block if we need to invalidate this instruction when EP escapes
-    if level == 0 && !jit.escapes_ep() && !jit.at_compile_target() {
+    if level == 0 && jit.can_assume_on_stack_env() && !jit.at_compile_target() {
         return jit.defer_compilation(asm);
     }
 
diff --git a/yjit/src/cruby.rs b/yjit/src/cruby.rs
index 22bc247959..5799ac70c1 100644
--- a/yjit/src/cruby.rs
+++ b/yjit/src/cruby.rs
@@ -606,6 +606,13 @@ impl From<VALUE> for u16 {
     }
 }
 
+/// Check whether a control frame has an escaped environment
+pub unsafe fn cfp_env_has_escaped(cfp: *mut rb_control_frame_struct) -> bool {
+    use crate::utils::IntoUsize;
+    let ep = get_cfp_ep(cfp);
+    0 != ep.offset(VM_ENV_DATA_INDEX_FLAGS as isize).read().0 & VM_ENV_FLAG_ESCAPED.as_usize()
+}
+
 /// Produce a Ruby string from a Rust string slice
 pub fn rust_str_to_ruby(str: &str) -> VALUE {
     unsafe { rb_utf8_str_new(str.as_ptr() as *const _, str.len() as i64) }
diff --git a/yjit/src/invariants.rs b/yjit/src/invariants.rs
index d468cfebd9..9c7392366a 100644
--- a/yjit/src/invariants.rs
+++ b/yjit/src/invariants.rs
@@ -168,8 +168,8 @@ pub fn track_no_ep_escape_assumption(uninit_block: BlockRef, iseq: IseqPtr) {
         .insert(uninit_block);
 }
 
-/// Returns true if a given ISEQ has previously escaped an environment.
-pub fn iseq_escapes_ep(iseq: IseqPtr) -> bool {
+/// Returns true if a given ISEQ has escaped an environment since YJIT boot.
+pub fn seen_escaped_env(iseq: IseqPtr) -> bool {
     Invariants::get_instance()
         .no_ep_escape_iseqs
         .get(&iseq)
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index c0b82fe70b..3240a9b405 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -279,6 +279,7 @@ pub const DEFAULT_COUNTERS: &'static [Counter] = &[
     Counter::compiled_blockid_count,
     Counter::compiled_block_count,
     Counter::deleted_defer_block_count,
+    Counter::exceptional_entry_escaped_env,
     Counter::compiled_branch_count,
     Counter::compile_time_ns,
     Counter::compilation_failure,
@@ -631,6 +632,8 @@ make_counters! {
     iseq_stack_too_large,
     iseq_too_long,
 
+    exceptional_entry_escaped_env,
+
     temp_reg_opnd,
     temp_mem_opnd,
     temp_spill,
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
index c3d851fdbc..5eb15b0201 100644
--- a/yjit/src/yjit.rs
+++ b/yjit/src/yjit.rs
@@ -151,6 +151,13 @@ pub extern "C" fn rb_yjit_iseq_gen_entry_point(iseq: IseqPtr, ec: EcPtr, jit_exc
         return std::ptr::null();
     }
 
+    // In case of exceptional entry, reject escaped environment.
+    // This allows us to use the fact that new frames generally start with an on-stack environment.
+    if jit_exception && unsafe { cfp_env_has_escaped(get_ec_cfp(ec)) } {
+        incr_counter!(exceptional_entry_escaped_env);
+        return std::ptr::null();
+    }
+
     // If a custom call threshold was not specified on the command-line and
     // this is a large application (has very many ISEQs), switch to
     // using the call threshold for large applications after this entry point