From 63c9aeed978fcd0942238772c2e441b33e12e16a Mon Sep 17 00:00:00 2001
From: Tom de Bruijn <tom@tomdebruijn.com>
Date: Thu, 4 Jul 2024 15:26:03 +0200
Subject: [PATCH] Add Rake performance instrumentation

We've had requests from people to instrument Rake tasks for performance
over time. I remember giving out a variation of this gist from time to
time to make this work for people that wanted it:
https://gist.github.com/tombruijn/701a3c5e3f251fb5e3ba8f1f8e908887

Let's add it to AppSignal itself so it can be enabled with a config
option, rather than a brittle monkeypatch on top of our existing
instrumentation.

It's turned off by default so we don't start collecting a lot more
'requests' that count towards an organization's plan, increasing the
cost of AppSignal without notice. We can always decide to enable this in
the next major version if we think it's useful to be turned on by
default.

Part of #1135
---
 ...d-rake-task-performance-instrumentation.md |  6 ++
 lib/appsignal/config.rb                       |  4 ++
 lib/appsignal/integrations/rake.rb            | 40 ++++++++-----
 spec/integration/diagnose                     |  2 +-
 spec/lib/appsignal/config_spec.rb             |  1 +
 spec/lib/appsignal/hooks/rake_spec.rb         | 58 ++++++++++++++-----
 6 files changed, 83 insertions(+), 28 deletions(-)
 create mode 100644 .changesets/add-rake-task-performance-instrumentation.md

diff --git a/.changesets/add-rake-task-performance-instrumentation.md b/.changesets/add-rake-task-performance-instrumentation.md
new file mode 100644
index 000000000..8232aa6ce
--- /dev/null
+++ b/.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.
diff --git a/lib/appsignal/config.rb b/lib/appsignal/config.rb
index 41aec49ea..db6758092 100644
--- a/lib/appsignal/config.rb
+++ b/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
diff --git a/lib/appsignal/integrations/rake.rb b/lib/appsignal/integrations/rake.rb
index 30e52ea52..8d4b82246 100644
--- a/lib/appsignal/integrations/rake.rb
+++ b/lib/appsignal/integrations/rake.rb
@@ -4,24 +4,38 @@ module Appsignal
   module Integrations
     module RakeIntegration
       def execute(*args)
-        super
+        transaction =
+          if Appsignal.config[:enable_rake_performance_instrumentation]
+            _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)
+        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
+          Appsignal.stop("rake")
+        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
-        Appsignal.stop("rake")
-        raise error
       end
     end
   end
diff --git a/spec/integration/diagnose b/spec/integration/diagnose
index 19b446805..7acb6e7c5 160000
--- a/spec/integration/diagnose
+++ b/spec/integration/diagnose
@@ -1 +1 @@
-Subproject commit 19b44680559bf2eab25d4c1ce6a144a6d22f8349
+Subproject commit 7acb6e7c56a93ec1e0e1fd148407b3e7461a08ef
diff --git a/spec/lib/appsignal/config_spec.rb b/spec/lib/appsignal/config_spec.rb
index 88b3fc3ad..2b4670948 100644
--- a/spec/lib/appsignal/config_spec.rb
+++ b/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                => [],
diff --git a/spec/lib/appsignal/hooks/rake_spec.rb b/spec/lib/appsignal/hooks/rake_spec.rb
index 9f56183d6..3a6c1e7de 100644
--- a/spec/lib/appsignal/hooks/rake_spec.rb
+++ b/spec/lib/appsignal/hooks/rake_spec.rb
@@ -3,39 +3,69 @@
 describe Appsignal::Hooks::RakeHook do
   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 { start_agent }
+  around { |example| keep_transactions { example.run } }
 
   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
+          expect(Appsignal).to_not receive(:stop)
+        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
       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
+
+          # 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")
+        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
       end
     end
 
     context "with error" do
-      let(:error) { ExampleException }
       before do
-        task.enhance { raise error, "my error message" }
+        task.enhance { raise ExampleException, "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")
       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,7 +75,7 @@ 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