From 281b03c7f088eac6eae38605f69b0428879eb75b Mon Sep 17 00:00:00 2001
From: Peter Oxenham <peter.oxenham@thoughtworks.com>
Date: Fri, 4 Mar 2016 16:12:20 +1000
Subject: [PATCH] Don't add custom field before controller is called and ensure
 custom fields are added when an exception is thrown

---
 .../metal/instrumentation.rb                  | 39 ++++++++++---------
 logstasher.iml                                |  9 +++++
 spec/lib/logstasher/instrumentation_spec.rb   | 39 +++++++++++++++++++
 3 files changed, 68 insertions(+), 19 deletions(-)
 create mode 100644 logstasher.iml

diff --git a/lib/logstasher/rails_ext/action_controller/metal/instrumentation.rb b/lib/logstasher/rails_ext/action_controller/metal/instrumentation.rb
index 11779bf..2d3aed1 100644
--- a/lib/logstasher/rails_ext/action_controller/metal/instrumentation.rb
+++ b/lib/logstasher/rails_ext/action_controller/metal/instrumentation.rb
@@ -29,26 +29,27 @@ def process_action(*args)
           logstasher_add_custom_fields_to_request_context(LogStasher.request_context)
         end
 
-        if respond_to?(:logstasher_add_custom_fields_to_payload)
-          before_keys = raw_payload.keys.clone
-          logstasher_add_custom_fields_to_payload(raw_payload)
-          after_keys = raw_payload.keys
-          # Store all extra keys added to payload hash in payload itself. This is a thread safe way
-          LogStasher::CustomFields.add(*(after_keys - before_keys))
+        begin
+          result = super
+          payload[:status] = response.status
+          result
+        ensure
+          if respond_to?(:logstasher_add_custom_fields_to_payload)
+            before_keys = raw_payload.keys.clone
+            logstasher_add_custom_fields_to_payload(raw_payload)
+            after_keys = raw_payload.keys
+            # Store all extra keys added to payload hash in payload itself. This is a thread safe way
+            LogStasher::CustomFields.add(*(after_keys - before_keys))
+          end
+          append_info_to_payload(payload)
+          LogStasher.store.each do |key, value|
+            payload[key] = value
+          end
+
+          LogStasher.request_context.each do |key, value|
+            payload[key] = value
+          end
         end
-
-        result = super
-
-        payload[:status] = response.status
-        append_info_to_payload(payload)
-        LogStasher.store.each do |key, value|
-          payload[key] = value
-        end
-
-        LogStasher.request_context.each do |key, value|
-          payload[key] = value
-        end
-        result
       end
     end
     alias logstasher_process_action process_action
diff --git a/logstasher.iml b/logstasher.iml
new file mode 100644
index 0000000..0ae9b6e
--- /dev/null
+++ b/logstasher.iml
@@ -0,0 +1,9 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="RUBY_MODULE" version="4">
+  <component name="NewModuleRootManager" inherit-compiler-output="true">
+    <exclude-output />
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="jdk" jdkName="rbenv: 2.2.3" jdkType="RUBY_SDK" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>
\ No newline at end of file
diff --git a/spec/lib/logstasher/instrumentation_spec.rb b/spec/lib/logstasher/instrumentation_spec.rb
index 328fa95..c9e83bc 100644
--- a/spec/lib/logstasher/instrumentation_spec.rb
+++ b/spec/lib/logstasher/instrumentation_spec.rb
@@ -70,5 +70,44 @@ def subject.index(*_args)
         end
       end
     end
+
+    context "payload has custom fields defined" do
+      before :each do
+        LogStasher.add_custom_fields do |fields|
+          fields[:some_field] = 'value'
+        end
+
+        ActiveSupport::Notifications.subscribe('process_action.action_controller') do |_, _, _, _, payload|
+          @payload = payload
+        end
+      end
+
+      it "should retain the value in the request context" do
+        subject.process_action(:index)
+      end
+
+      context 'exception thrown by controller' do
+        before :each do
+          def subject.index_with_exception(*args)
+            fail 'Exception'
+          end
+        end
+
+        it 'should retain the value in the request context' do
+          expect { subject.process_action(:index_with_exception) }.to raise_exception('Exception')
+        end
+      end
+
+      after :each do
+        expect(@payload[:some_field]).to eq('value')
+
+        ActionController::Metal.class_eval do
+          undef logstasher_add_custom_fields_to_payload
+        end
+        ActionController::Base.class_eval do
+          undef logstasher_add_custom_fields_to_payload
+        end
+      end
+    end
   end
 end