From 28dc515d33e3de94ef256d0e3b68b4aebd11510c Mon Sep 17 00:00:00 2001
From: Samuel Cochran <sj26@sj26.com>
Date: Tue, 7 Mar 2023 10:51:56 +1100
Subject: [PATCH] Allow disabling per-process metrics

---
 lib/sidekiq/cloudwatchmetrics.rb       | 33 ++++++++++++++------------
 spec/sidekiq/cloudwatchmetrics_spec.rb | 23 ++++++++++++++++++
 2 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/lib/sidekiq/cloudwatchmetrics.rb b/lib/sidekiq/cloudwatchmetrics.rb
index b4739f3..6389328 100644
--- a/lib/sidekiq/cloudwatchmetrics.rb
+++ b/lib/sidekiq/cloudwatchmetrics.rb
@@ -45,12 +45,13 @@ class Publisher
 
     INTERVAL = 60 # seconds
 
-    def initialize(config: Sidekiq, client: Aws::CloudWatch::Client.new, namespace: "Sidekiq", additional_dimensions: {})
+    def initialize(config: Sidekiq, client: Aws::CloudWatch::Client.new, namespace: "Sidekiq", process_metrics: true, additional_dimensions: {})
       # Sidekiq 6.5+ requires @config, which defaults to the top-level
       # `Sidekiq` module, but can be overridden when running multiple Sidekiqs.
       @config = config
       @client = client
       @namespace = namespace
+      @process_metrics = process_metrics
       @additional_dimensions = additional_dimensions.map { |k, v| {name: k.to_s, value: v.to_s} }
     end
 
@@ -163,23 +164,25 @@ def publish
         }
       end
 
-      processes.each do |process|
-        process_utilization = process["busy"] / process["concurrency"].to_f * 100.0
+      if @process_metrics
+        processes.each do |process|
+          process_utilization = process["busy"] / process["concurrency"].to_f * 100.0
 
-        unless process_utilization.nan?
-          process_dimensions = [{name: "Hostname", value: process["hostname"]}]
+          unless process_utilization.nan?
+            process_dimensions = [{name: "Hostname", value: process["hostname"]}]
 
-          if process["tag"]
-            process_dimensions << {name: "Tag", value: process["tag"]}
-          end
+            if process["tag"]
+              process_dimensions << {name: "Tag", value: process["tag"]}
+            end
 
-          metrics << {
-            metric_name: "Utilization",
-            dimensions: process_dimensions,
-            timestamp: now,
-            value: process_utilization,
-            unit: "Percent",
-          }
+            metrics << {
+              metric_name: "Utilization",
+              dimensions: process_dimensions,
+              timestamp: now,
+              value: process_utilization,
+              unit: "Percent",
+            }
+          end
         end
       end
 
diff --git a/spec/sidekiq/cloudwatchmetrics_spec.rb b/spec/sidekiq/cloudwatchmetrics_spec.rb
index c1d52d5..057cb88 100644
--- a/spec/sidekiq/cloudwatchmetrics_spec.rb
+++ b/spec/sidekiq/cloudwatchmetrics_spec.rb
@@ -353,6 +353,29 @@
           end
         end
       end
+
+      context "when per process metrics are disabled" do
+        subject(:publisher) { Sidekiq::CloudWatchMetrics::Publisher.new(client: client, process_metrics: false) }
+
+        it "only publishes a single Utilization metric" do
+          Timecop.freeze(now = Time.now) do
+            publisher.publish
+
+            expect(client).to have_received(:put_metric_data) { |metrics|
+              utilization_data = metrics[:metric_data].select { |data| data[:metric_name] == "Utilization" }
+
+              expect(utilization_data).to contain_exactly(
+                {
+                  metric_name: "Utilization",
+                  timestamp: now,
+                  value: 30.0,
+                  unit: "Percent",
+                },
+              )
+            }
+          end
+        end
+      end
     end
 
     describe "#stop" do