From 104946b852e128b70013c216cd703e7e830b4981 Mon Sep 17 00:00:00 2001
From: Stan Lo <stan001212@gmail.com>
Date: Thu, 21 Mar 2024 15:17:23 +0000
Subject: [PATCH] Support rescue events in the exception tracer

This feature is only activated in Ruby 3.3 and later. It allows the
tracer to print trace when exceptions are rescued.

Co-authored-by: Kaan Ozkan <kaan.ozkan@shopify.com>
---
 README.md                            |  3 ++-
 lib/tracer/exception_tracer.rb       | 32 +++++++++++++++++++++-------
 test/tracer/exception_tracer_test.rb | 30 ++++++++++++++++++++++----
 test/tracer/helper_test.rb           | 10 ++++++++-
 4 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index f8b1c16..7e9d002 100644
--- a/README.md
+++ b/README.md
@@ -139,7 +139,8 @@ rescue StandardError
   nil
 end
 
- #depth:1  #<RuntimeError: boom> at test.rb:4
+#depth:0  #<RuntimeError: boom> raised at test.rb:4
+#depth:1  #<RuntimeError: boom> rescued at test.rb:6
 ```
 
 #### CallTracer
diff --git a/lib/tracer/exception_tracer.rb b/lib/tracer/exception_tracer.rb
index 62896ee..83a10c0 100644
--- a/lib/tracer/exception_tracer.rb
+++ b/lib/tracer/exception_tracer.rb
@@ -4,16 +4,32 @@
 
 class ExceptionTracer < Tracer::Base
   def setup_tp
-    TracePoint.new(:raise) do |tp|
-      next if skip?(tp)
+    if RUBY_VERSION >= "3.3.0"
+      TracePoint.new(:raise, :rescue) do |tp|
+        next if skip?(tp)
 
-      exc = tp.raised_exception
+        exc = tp.raised_exception
 
-      out tp,
-          " #{colorize_magenta(exc.inspect)}",
-          depth: caller.size - (1 + @depth_offset)
-    rescue Exception => e
-      p e
+        action = tp.event == :raise ? "raised" : "rescued"
+
+        out tp,
+            " #{colorize_magenta(exc.inspect)} #{action}",
+            depth: caller.size - (1 + @depth_offset)
+      rescue Exception => e
+        p e
+      end
+    else
+      TracePoint.new(:raise) do |tp|
+        next if skip?(tp)
+
+        exc = tp.raised_exception
+
+        out tp,
+            " #{colorize_magenta(exc.inspect)} raised",
+            depth: caller.size - (1 + @depth_offset)
+      rescue Exception => e
+        p e
+      end
     end
   end
 
diff --git a/test/tracer/exception_tracer_test.rb b/test/tracer/exception_tracer_test.rb
index 1f40b05..47faf79 100644
--- a/test/tracer/exception_tracer_test.rb
+++ b/test/tracer/exception_tracer_test.rb
@@ -16,26 +16,48 @@ def test_exception_tracer_traces_exceptions
       file = write_file("foo.rb", <<~RUBY)
         ExceptionTracer.new.start
 
-        raise "boom" rescue nil
+        begin
+          raise "boom"
+        rescue
+        end
       RUBY
 
       out, err = execute_file(file)
 
+      expected_traces = [
+        /^#depth:0  #<RuntimeError: boom> raised at .*foo.rb:4/
+      ]
+
+      if RUBY_VERSION >= "3.3.0"
+        expected_traces << /^#depth:1  #<RuntimeError: boom> rescued at .*foo.rb:5/
+      end
+
       assert_empty(err)
-      assert_traces([/^#depth:0  #<RuntimeError: boom> at .*foo.rb:3/], out)
+      assert_traces(expected_traces, out)
     end
 
     def test_exception_tracer_with_header
       file = write_file("foo.rb", <<~RUBY)
         ExceptionTracer.new(header: "tracer-1").start
 
-        raise "boom" rescue nil
+        begin
+          raise "boom"
+        rescue
+        end
       RUBY
 
       out, err = execute_file(file)
 
+      expected_traces = [
+        /^tracer-1 #depth:0  #<RuntimeError: boom> raised at .*foo.rb:4/
+      ]
+
+      if RUBY_VERSION >= "3.3.0"
+        expected_traces << /^tracer-1 #depth:1  #<RuntimeError: boom> rescued at .*foo.rb:5/
+      end
+
       assert_empty(err)
-      assert_traces([/^tracer-1 #depth:0  #<RuntimeError: boom> at .*foo.rb:3/], out)
+      assert_traces(expected_traces, out)
     end
   end
 end
diff --git a/test/tracer/helper_test.rb b/test/tracer/helper_test.rb
index d601dd3..73a87f6 100644
--- a/test/tracer/helper_test.rb
+++ b/test/tracer/helper_test.rb
@@ -13,8 +13,16 @@ def test_trace_exception
 
       out, err = execute_file(file)
 
+      expected_traces = [
+        /^#depth:1  #<RuntimeError: boom> raised at .*foo.rb:4/
+      ]
+
+      if RUBY_VERSION >= "3.3.0"
+        expected_traces << /^#depth:2  #<RuntimeError: boom> rescued at .*foo.rb:4/
+      end
+
       assert_empty(err)
-      assert_traces([/#depth:1  #<RuntimeError: boom> at .*foo.rb:4/], out)
+      assert_traces(expected_traces, out)
     end
 
     def test_trace_call