Don't add custom field before controller is called and ensure custom …

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

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>

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