From 9873af0b1a343dff6d1a8af4c813aa2c9ecc47d5 Mon Sep 17 00:00:00 2001
From: Koichi Sasada <ko1@atdot.net>
Date: Thu, 9 Dec 2021 03:50:17 +0900
Subject: [PATCH] `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>
---
 test/ruby/test_settracefunc.rb | 24 ++++++++++++++++++++++++
 trace_point.rb                 | 16 ++++++++++++++++
 vm_trace.c                     | 19 +++++++++++++++++++
 3 files changed, 59 insertions(+)

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
@@ -135,6 +135,22 @@ class TracePoint
     Primitive.tracepoint_trace_s(events)
   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 */