From bea7450d8f067539521c00c7c339b0837a83e42b Mon Sep 17 00:00:00 2001
From: Adam Cooke <adam@adam.ac>
Date: Sat, 24 Feb 2024 11:19:48 +0000
Subject: [PATCH] feat: add prometheus metrics to worker

---
 app/lib/worker/process.rb          | 33 +++++++++++++++++++++++++++---
 app/util/has_prometheus_metrics.rb | 12 +++++++++++
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/app/lib/worker/process.rb b/app/lib/worker/process.rb
index 1e17dc3..11167e5 100644
--- a/app/lib/worker/process.rb
+++ b/app/lib/worker/process.rb
@@ -139,7 +139,15 @@ module Worker
         JOBS.each do |job_class|
           capture_errors do
             job = job_class.new(logger: logger)
-            job.call
+
+            time = Benchmark.realtime { job.call }
+
+            observe_prometheus_histogram :postal_worker_job_runtime,
+                                         time,
+                                         labels: {
+                                          thread: thread,
+                                          job: job_class.to_s.split("::").last
+                                         }
 
             if job.work_completed?
               completed_work += 1
@@ -222,7 +230,18 @@ module Worker
 
         logger.info "running task"
 
-        capture_errors { task.new(logger: logger).call }
+        time = 0
+        capture_errors do
+          time = Benchmark.realtime do
+            task.new(logger: logger).call
+          end
+
+          observe_prometheus_histogram :postal_worker_task_runtime,
+                                       time,
+                                       labels: {
+                                        task: task.to_s.split("::").last
+                                       }
+        end
 
         next_run_after = task.next_run_after
         logger.info "scheduling task to next run at #{next_run_after}"
@@ -254,12 +273,20 @@ module Worker
 
     def setup_prometheus
       register_prometheus_counter :postal_worker_job_executions,
-                                  docstring: "The number of jobs worked by a worker",
+                                  docstring: "The number of jobs worked by a worker where work was completed",
                                   labels: [:thread, :job]
 
+      register_prometheus_histogram :postal_worker_job_runtime,
+                                    docstring: "The time taken to process jobs",
+                                    labels: [:thread, :job]
+
       register_prometheus_counter :postal_worker_errors,
                                   docstring: "The number of errors encountered while processing jobs",
                                   labels: [:error]
+
+      register_prometheus_histogram :postal_worker_task_runtime,
+                                    docstring: "The time taken to process tasks",
+                                    labels: [:task]
     end
 
   end
diff --git a/app/util/has_prometheus_metrics.rb b/app/util/has_prometheus_metrics.rb
index ca71731..bd2224e 100644
--- a/app/util/has_prometheus_metrics.rb
+++ b/app/util/has_prometheus_metrics.rb
@@ -7,6 +7,11 @@ module HasPrometheusMetrics
     registry.register(counter)
   end
 
+  def register_prometheus_histogram(name, **kwargs)
+    histogram = Prometheus::Client::Histogram.new(name, **kwargs)
+    registry.register(histogram)
+  end
+
   def increment_prometheus_counter(name, labels: {})
     counter = registry.get(name)
     return if counter.nil?
@@ -14,6 +19,13 @@ module HasPrometheusMetrics
     counter.increment(labels: labels)
   end
 
+  def observe_prometheus_histogram(name, time, labels: {})
+    histogram = registry.get(name)
+    return if histogram.nil?
+
+    histogram.observe(time, labels: labels)
+  end
+
   private
 
   def registry