From 9094fdf04e767e43c57b2b8e1b4c2bf4bfff9e07 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 1/2] 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 f7a5c4c..da0a3ee 100644
--- a/lib/logstasher/rails_ext/action_controller/metal/instrumentation.rb
+++ b/lib/logstasher/rails_ext/action_controller/metal/instrumentation.rb
@@ -23,26 +23,27 @@ def process_action(*args)
           logstasher_add_custom_fields_to_request_context(LogStasher.request_context)
         end
 
-        if self.respond_to?(:logtasher_add_custom_fields_to_payload)
-          before_keys = raw_payload.keys.clone
-          logtasher_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.custom_fields += after_keys - before_keys
+        begin
+          result = super
+          payload[:status] = response.status
+          result
+        ensure
+          if self.respond_to?(:logtasher_add_custom_fields_to_payload)
+            before_keys = raw_payload.keys.clone
+            logtasher_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.custom_fields += 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 74cbcad..57920e0 100644
--- a/spec/lib/logstasher/instrumentation_spec.rb
+++ b/spec/lib/logstasher/instrumentation_spec.rb
@@ -66,5 +66,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 logtasher_add_custom_fields_to_payload
+        end
+        ActionController::Base.class_eval do
+          undef logtasher_add_custom_fields_to_payload
+        end
+      end
+    end
   end
 end

From 7910793ea3e6bf7edb625fc34c520144f7be24b2 Mon Sep 17 00:00:00 2001
From: Peter Oxenham <peter.oxenham@thoughtworks.com>
Date: Fri, 8 Jul 2016 15:24:58 +1000
Subject: [PATCH 2/2] Fixed custom field rails extension

---
 .../rails_ext/action_controller/metal/instrumentation.rb        | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/logstasher/rails_ext/action_controller/metal/instrumentation.rb b/lib/logstasher/rails_ext/action_controller/metal/instrumentation.rb
index 7507855..24c363c 100644
--- a/lib/logstasher/rails_ext/action_controller/metal/instrumentation.rb
+++ b/lib/logstasher/rails_ext/action_controller/metal/instrumentation.rb
@@ -33,7 +33,7 @@ def process_action(*args)
             logtasher_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.custom_fields += after_keys - before_keys
+            LogStasher::CustomFields.custom_fields += after_keys - before_keys
           end
           append_info_to_payload(payload)
           LogStasher.store.each do |key, value|