Merge pull request #1156 from appsignal/rake-performance-instrumentation

Add Rake performance instrumentation

.changesets/add-rake-task-performance-instrumentation.md

@@ -0,0 +1,6 @@
+---
+bump: minor
+type: add
+---
+
+Add Rake task performance instrumentation. Configure the `enable_rake_performance_instrumentation` option to `true` to enable Rake task instrumentation for both error and performance monitoring. To ignore specific Rake tasks, configure `ignore_actions` to include the name of the Rake task.

lib/appsignal/config.rb

@@ -23,6 +23,7 @@ class Config
       :enable_gvl_global_timer => true,
       :enable_gvl_waiting_threads => true,
       :enable_rails_error_reporter => true,
+      :enable_rake_performance_instrumentation => false,
       :endpoint => "https://push.appsignal.com",
       :files_world_accessible => true,
       :filter_metadata => [],
@@ -83,6 +84,8 @@ class Config
       "APPSIGNAL_ENABLE_GVL_GLOBAL_TIMER" => :enable_gvl_global_timer,
       "APPSIGNAL_ENABLE_GVL_WAITING_THREADS" => :enable_gvl_waiting_threads,
       "APPSIGNAL_ENABLE_RAILS_ERROR_REPORTER" => :enable_rails_error_reporter,
+      "APPSIGNAL_ENABLE_RAKE_PERFORMANCE_INSTRUMENTATION" =>
+        :enable_rake_performance_instrumentation,
       "APPSIGNAL_FILES_WORLD_ACCESSIBLE" => :files_world_accessible,
       "APPSIGNAL_FILTER_METADATA" => :filter_metadata,
       "APPSIGNAL_FILTER_PARAMETERS" => :filter_parameters,
@@ -150,6 +153,7 @@ class Config
       APPSIGNAL_ENABLE_GVL_GLOBAL_TIMER
       APPSIGNAL_ENABLE_GVL_WAITING_THREADS
       APPSIGNAL_ENABLE_RAILS_ERROR_REPORTER
+      APPSIGNAL_ENABLE_RAKE_PERFORMANCE_INSTRUMENTATION
       APPSIGNAL_FILES_WORLD_ACCESSIBLE
       APPSIGNAL_INSTRUMENT_HTTP_RB
       APPSIGNAL_INSTRUMENT_NET_HTTP

lib/appsignal/integrations/rake.rb

@@ -2,26 +2,60 @@
 
 module Appsignal
   module Integrations
+    # @api private
     module RakeIntegration
       def execute(*args)
-        super
+        transaction =
+          if Appsignal.config[:enable_rake_performance_instrumentation]
+            Appsignal::Integrations::RakeIntegrationHelper.register_at_exit_hook
+            _appsignal_create_transaction
+          end
+
+        Appsignal.instrument "task.rake" do
+          super
+        end
       rescue Exception => error # rubocop:disable Lint/RescueException
-        # Format given arguments and cast to hash if possible
-        params, _ = args
-        params = params.to_hash if params.respond_to?(:to_hash)
+        Appsignal::Integrations::RakeIntegrationHelper.register_at_exit_hook
+        transaction ||= _appsignal_create_transaction
+        transaction.set_error(error)
+        raise error
+      ensure
+        if transaction
+          # Format given arguments and cast to hash if possible
+          params, _ = args
+          params = params.to_hash if params.respond_to?(:to_hash)
+          transaction.set_params_if_nil(params)
+          transaction.set_action(name)
+          transaction.complete
+        end
+      end
 
-        transaction = Appsignal::Transaction.create(
+      private
+
+      def _appsignal_create_transaction
+        Appsignal::Transaction.create(
           SecureRandom.uuid,
           Appsignal::Transaction::BACKGROUND_JOB,
-          Appsignal::Transaction::GenericRequest.new(
-            :params => params
-          )
+          Appsignal::Transaction::GenericRequest.new({})
         )
-        transaction.set_action(name)
-        transaction.set_error(error)
-        transaction.complete
+      end
+    end
+
+    # @api private
+    module RakeIntegrationHelper
+      # Register an `at_exit` hook when a task is executed. This will stop
+      # AppSignal when _all_ tasks are executed and Rake exits.
+      def self.register_at_exit_hook
+        return if @register_at_exit_hook
+
+        Kernel.at_exit(&method(:at_exit_hook))
+
+        @register_at_exit_hook = true
+      end
+
+      # The at_exit hook itself
+      def self.at_exit_hook
         Appsignal.stop("rake")
-        raise error
       end
     end
   end

spec/lib/appsignal/config_spec.rb

@@ -165,6 +165,7 @@
         :enable_statsd                  => true,
         :enable_nginx_metrics           => false,
         :enable_rails_error_reporter    => true,
+        :enable_rake_performance_instrumentation => false,
         :endpoint                       => "https://push.appsignal.com",
         :files_world_accessible         => true,
         :filter_metadata                => [],

spec/lib/appsignal/hooks/rake_spec.rb

@@ -1,41 +1,89 @@
 require "rake"
 
 describe Appsignal::Hooks::RakeHook do
+  let(:helper) { Appsignal::Integrations::RakeIntegrationHelper }
   let(:task) { Rake::Task.new("task:name", Rake::Application.new) }
   let(:arguments) { Rake::TaskArguments.new(["foo"], ["bar"]) }
-  let(:generic_request) { Appsignal::Transaction::GenericRequest.new({}) }
-  before(:context) { start_agent }
+  before do
+    start_agent
+    allow(Kernel).to receive(:at_exit)
+  end
+  around { |example| keep_transactions { example.run } }
+  after do
+    if helper.instance_variable_defined?(:@register_at_exit_hook)
+      helper.remove_instance_variable(:@register_at_exit_hook)
+    end
+  end
+
+  def expect_to_not_have_registered_at_exit_hook
+    expect(Kernel).to_not have_received(:at_exit)
+  end
+
+  def expect_to_have_registered_at_exit_hook
+    expect(Kernel).to have_received(:at_exit)
+  end
 
   describe "#execute" do
     context "without error" do
-      before { expect(Appsignal).to_not receive(:stop) }
-
       def perform
         task.execute(arguments)
       end
 
-      it "creates no transaction" do
-        expect { perform }.to_not(change { created_transactions.count })
+      context "with :enable_rake_performance_instrumentation == false" do
+        before do
+          Appsignal.config[:enable_rake_performance_instrumentation] = false
+        end
+
+        it "creates no transaction" do
+          expect { perform }.to_not(change { created_transactions.count })
+        end
+
+        it "calls the original task" do
+          expect(perform).to eq([])
+        end
+
+        it "does not register an at_exit hook" do
+          perform
+          expect_to_not_have_registered_at_exit_hook
+        end
       end
 
-      it "calls the original task" do
-        expect(perform).to eq([])
+      context "with :enable_rake_performance_instrumentation == true" do
+        before do
+          Appsignal.config[:enable_rake_performance_instrumentation] = true
+        end
+
+        it "creates a transaction" do
+          expect { perform }.to(change { created_transactions.count }.by(1))
+
+          transaction = last_transaction
+          expect(transaction).to have_id
+          expect(transaction).to have_namespace(Appsignal::Transaction::BACKGROUND_JOB)
+          expect(transaction).to have_action("task:name")
+          expect(transaction).to_not have_error
+          expect(transaction).to include_params("foo" => "bar")
+          expect(transaction).to include_event("name" => "task.rake")
+          expect(transaction).to be_completed
+        end
+
+        it "calls the original task" do
+          expect(perform).to eq([])
+        end
+
+        it "registers an at_exit hook" do
+          perform
+          expect_to_have_registered_at_exit_hook
+        end
       end
     end
 
     context "with error" do
-      let(:error) { ExampleException }
       before do
-        task.enhance { raise error, "my error message" }
-        # We don't call `and_call_original` here as we don't want AppSignal to
-        # stop and start for every spec.
-        expect(Appsignal).to receive(:stop).with("rake")
+        task.enhance { raise ExampleException, "error message" }
       end
 
       def perform
-        keep_transactions do
-          expect { task.execute(arguments) }.to raise_error(error)
-        end
+        expect { task.execute(arguments) }.to raise_error(ExampleException, "error message")
       end
 
       it "creates a background job transaction" do
@@ -45,11 +93,16 @@ def perform
         expect(transaction).to have_id
         expect(transaction).to have_namespace(Appsignal::Transaction::BACKGROUND_JOB)
         expect(transaction).to have_action("task:name")
-        expect(transaction).to have_error("ExampleException", "my error message")
+        expect(transaction).to have_error("ExampleException", "error message")
         expect(transaction).to include_params("foo" => "bar")
         expect(transaction).to be_completed
       end
 
+      it "registers an at_exit hook" do
+        perform
+        expect_to_have_registered_at_exit_hook
+      end
+
       context "when first argument is not a `Rake::TaskArguments`" do
         let(:arguments) { nil }
 
@@ -62,3 +115,33 @@ def perform
     end
   end
 end
+
+describe "Appsignal::Integrations::RakeIntegrationHelper" do
+  let(:helper) { Appsignal::Integrations::RakeIntegrationHelper }
+  describe ".register_at_exit_hook" do
+    before do
+      start_agent
+      allow(Appsignal).to receive(:stop)
+    end
+
+    it "registers the at_exit hook only once" do
+      allow(Kernel).to receive(:at_exit)
+      helper.register_at_exit_hook
+      helper.register_at_exit_hook
+      expect(Kernel).to have_received(:at_exit).once
+    end
+  end
+
+  describe ".at_exit_hook" do
+    let(:helper) { Appsignal::Integrations::RakeIntegrationHelper }
+    before do
+      start_agent
+      allow(Appsignal).to receive(:stop)
+    end
+
+    it "calls Appsignal.stop" do
+      helper.at_exit_hook
+      expect(Appsignal).to have_received(:stop).with("rake")
+    end
+  end
+end