`TracePoint.allow_reentry`

In general, while TracePoint callback is running,
other registerred callbacks are not called to avoid
confusion by reentrace.
This method allow the reentrace. This method should be
used carefully, otherwize the callback can be easily called
infinitely.

[Feature #15912]

Co-authored-by: Jean Boussier <jean.boussier@gmail.com>

3 files changed, 59 insertions, 0 deletions

diff --git a/test/ruby/test_settracefunc.rb b/test/ruby/test_settracefunc.rb
index d737c84af5..f623d87506 100644
--- a/test/ruby/test_settracefunc.rb
+++ b/test/ruby/test_settracefunc.rb
@@ -2575,4 +2575,28 @@ CODE
     }
     assert_equal [__LINE__ - 5, __LINE__ - 4, __LINE__ - 3], lines, 'Bug #17868'
   end
+
+  def test_allow_reentry
+    event_lines = []
+    _l1 = _l2 = _l3 = _l4 = nil
+    TracePoint.new(:line) do |tp|
+      next unless target_thread?
+
+      event_lines << tp.lineno
+      next if (__LINE__ + 2 .. __LINE__ + 4).cover?(tp.lineno)
+      TracePoint.allow_reentry do
+        _a = 1; _l3 = __LINE__
+        _b = 2; _l4 = __LINE__
+      end
+    end.enable do
+      _c = 3; _l1 = __LINE__
+      _d = 4; _l2 = __LINE__
+    end
+
+    assert_equal [_l1, _l3, _l4, _l2, _l3, _l4], event_lines
+
+    assert_raise RuntimeError do
+      TracePoint.allow_reentry{}
+    end
+  end
 end
diff --git a/trace_point.rb b/trace_point.rb
index db34e1b68d..85ebac9aa7 100644
--- a/trace_point.rb
+++ b/trace_point.rb
@@ -136,6 +136,22 @@ class TracePoint
   end
 
   # call-seq:
+  #   TracePoint.allow_reentry
+  #
+  # In general, while a TracePoint callback is running,
+  # other registered callbacks are not called to avoid
+  # confusion by reentrance.
+  # This method allows the reentrance in a given block.
+  # This method should be used carefully, otherwise the callback
+  # can be easily called infinitely.
+  #
+  # If this method is called when the reentrance is already allowed,
+  # it raises a RuntimeError.
+  def self.allow_reentry
+    Primitive.tracepoint_allow_reentry
+  end
+
+  # call-seq:
   #    trace.enable(target: nil, target_line: nil, target_thread: nil)    -> true or false
   #    trace.enable(target: nil, target_line: nil, target_thread: nil) { block }  -> obj
   #
diff --git a/vm_trace.c b/vm_trace.c
index 3a0f7a7aad..f1c0bb2598 100644
--- a/vm_trace.c
+++ b/vm_trace.c
@@ -1517,6 +1517,25 @@ tracepoint_stat_s(rb_execution_context_t *ec, VALUE self)
     return stat;
 }
 
+static VALUE
+disallow_reentry(VALUE val)
+{
+    rb_trace_arg_t *arg = (rb_trace_arg_t *)val;
+    rb_execution_context_t *ec = GET_EC();
+    if (ec->trace_arg != NULL) rb_bug("should be NULL, but %p", (void *)ec->trace_arg);
+    ec->trace_arg = arg;
+    return Qnil;
+}
+
+static VALUE
+tracepoint_allow_reentry(rb_execution_context_t *ec, VALUE self)
+{
+    const rb_trace_arg_t *arg = ec->trace_arg;
+    if (arg == NULL) rb_raise(rb_eRuntimeError, "No need to allow reentrance.");
+    ec->trace_arg = NULL;
+    return rb_ensure(rb_yield, Qnil, disallow_reentry, (VALUE)arg);
+}
+
 #include "trace_point.rbinc"
 
 /* This function is called from inits.c */