From 31b77207267b77126e9c4fa59f8350acabea023b Mon Sep 17 00:00:00 2001
From: Marco Costa <marco.costa@datadoghq.com>
Date: Fri, 22 Nov 2024 11:58:30 -0800
Subject: [PATCH] Add exception to instrumentation payload

---
 lib/karafka/core/monitoring/notifications.rb  | 39 ++++++++++---------
 .../core/monitoring/notifications_spec.rb     | 26 +++++++++++++
 2 files changed, 47 insertions(+), 18 deletions(-)

diff --git a/lib/karafka/core/monitoring/notifications.rb b/lib/karafka/core/monitoring/notifications.rb
index 2d7bf00..f6ae54d 100644
--- a/lib/karafka/core/monitoring/notifications.rb
+++ b/lib/karafka/core/monitoring/notifications.rb
@@ -117,27 +117,30 @@ def instrument(event_id, payload = EMPTY_HASH)
             return yield if assigned_listeners.empty?
 
             start = monotonic_now
-            result = yield
-            time = monotonic_now - start
-          elsif assigned_listeners.empty?
-            # Skip measuring or doing anything if no one listening
-            return
+            begin
+              yield
+            rescue Exception => e
+              payload[:exception] = e
+              raise e
+            ensure
+              time = monotonic_now - start
+            end
           end
-
-          event = Event.new(
-            event_id,
-            time ? payload.merge(time: time) : payload
-          )
-
-          assigned_listeners.each do |listener|
-            if listener.is_a?(Proc)
-              listener.call(event)
-            else
-              listener.send(@events_methods_map[event_id], event)
+        ensure
+          if assigned_listeners && !assigned_listeners.empty?
+            event = Event.new(
+              event_id,
+              time ? payload.merge(time: time) : payload
+            )
+
+            assigned_listeners.each do |listener|
+              if listener.is_a?(Proc)
+                listener.call(event)
+              else
+                listener.send(@events_methods_map[event_id], event)
+              end
             end
           end
-
-          result
         end
       end
     end
diff --git a/spec/lib/karafka/core/monitoring/notifications_spec.rb b/spec/lib/karafka/core/monitoring/notifications_spec.rb
index d0198c5..8c51e78 100644
--- a/spec/lib/karafka/core/monitoring/notifications_spec.rb
+++ b/spec/lib/karafka/core/monitoring/notifications_spec.rb
@@ -28,6 +28,32 @@
         expect { notifications.instrument('na') }.to raise_error(expected_error)
       end
     end
+
+    context 'when the instrumented code errors' do
+      let(:result) { raise ArgumentError, "error message" }
+
+      it 'expect to raise the error' do
+        expect { instrumentation }.to raise_error(ArgumentError, "error message")
+      end
+
+      context 'with a subscriber' do
+        let(:subscriber_block) { double('block') }
+
+        before do
+          allow(subscriber_block).to receive(:call)
+          notifications.subscribe(event_name, &subscriber_block.method(:call))
+        end
+
+        it 'expect to provide exception information in event' do
+          expect(subscriber_block).to receive(:call) do |event|
+            expect(event[:exception]).to be_a(ArgumentError)
+            expect(event[:exception].message).to eq("error message")
+          end
+
+          expect { instrumentation }.to raise_error(ArgumentError)
+        end
+      end
+    end
   end
 
   describe '#subscribe' do