From 7cbf78a45870f96811965f098ee954708999807d Mon Sep 17 00:00:00 2001
From: wint3rmute <mateusz.baczek1998@gmail.com>
Date: Sun, 19 Mar 2023 16:28:10 +0100
Subject: [PATCH] Collecting num of requests and handling time from each Kemal
 route

---
 src/invidious/helpers/helpers.cr    |  8 -----
 src/invidious/metrics.cr            | 48 +++++++++++++++++++++++++++++
 src/invidious/routes/api/v1/misc.cr | 27 +++++++++++++++-
 src/invidious/routing.cr            |  1 +
 4 files changed, 75 insertions(+), 9 deletions(-)
 create mode 100644 src/invidious/metrics.cr

diff --git a/src/invidious/helpers/helpers.cr b/src/invidious/helpers/helpers.cr
index 34df6fe6..c3b53339 100644
--- a/src/invidious/helpers/helpers.cr
+++ b/src/invidious/helpers/helpers.cr
@@ -233,11 +233,3 @@ def proxy_file(response, env)
     IO.copy response.body_io, env.response
   end
 end
-
-def to_prometheus_metrics(metrics_struct : Hash(String, Number)) : String
-  return String.build do |str|
-    metrics_struct.each.each do |key, value|
-      str << key << " " << value << "\n"
-    end
-  end
-end
diff --git a/src/invidious/metrics.cr b/src/invidious/metrics.cr
new file mode 100644
index 00000000..8007ca7d
--- /dev/null
+++ b/src/invidious/metrics.cr
@@ -0,0 +1,48 @@
+# Module containing an optional Kemal handler, which can be used
+# to collect metrics about the usage of various Invidious routes.
+module Metrics
+  record MetricLabels, request_method : String, request_route : String, response_code : Int32
+
+  # Counts how many a given route was used
+  REQUEST_COUNTERS              = Hash(MetricLabels, Int64).new
+
+  # Counts how much time was used to handle requests to each route
+  REQUEST_DURATION_SECONDS_SUMS = Hash(MetricLabels, Float32).new
+
+  # The handler which will record metrics when registered in a Kemal application 
+  METRICS_COLLECTOR = RouteMetricsCollector.new(REQUEST_COUNTERS, REQUEST_DURATION_SECONDS_SUMS)
+
+  class RouteMetricsCollector < Kemal::Handler
+    def initialize(
+      @num_of_request_counters : Hash(MetricLabels, Int64),
+      @request_duration_seconds_sums : Hash(MetricLabels, Float32)
+    )
+    end
+
+    def call(context : HTTP::Server::Context)
+      request_handling_started = Time.utc
+      begin
+        call_next(context)
+      ensure
+        request_handling_finished = Time.utc
+        request_path = context.route.path
+        request_method = context.request.method
+        seconds_spent_handling = (request_handling_finished - request_handling_started).to_f
+        response_status = context.response.status_code.to_i
+
+        LOGGER.trace("Collecting metrics: handling #{request_method} #{request_path} took #{seconds_spent_handling}s and finished with status #{response_status}")
+        metric_key = MetricLabels.new request_path, request_method, response_status
+
+        unless @num_of_request_counters.has_key?(metric_key)
+          @num_of_request_counters[metric_key] = 0
+        end
+        @num_of_request_counters[metric_key] += 1
+
+        unless @request_duration_seconds_sums.has_key?(metric_key)
+          @request_duration_seconds_sums[metric_key] = 0.0
+        end
+        @request_duration_seconds_sums[metric_key] += seconds_spent_handling
+      end
+    end
+  end
+end
diff --git a/src/invidious/routes/api/v1/misc.cr b/src/invidious/routes/api/v1/misc.cr
index edc0f302..7159ff9d 100644
--- a/src/invidious/routes/api/v1/misc.cr
+++ b/src/invidious/routes/api/v1/misc.cr
@@ -1,3 +1,5 @@
+require "../../../metrics.cr"
+
 module Invidious::Routes::API::V1::Misc
   # Stats API endpoint for Invidious
   def self.stats(env)
@@ -12,7 +14,30 @@ module Invidious::Routes::API::V1::Misc
 
   def self.metrics(env)
     env.response.content_type = "text/plain"
-    return to_prometheus_metrics(Invidious::Jobs::StatisticsRefreshJob::STATISTICS_PROMETHEUS)
+
+    return String.build do |str|
+      Metrics::REQUEST_COUNTERS.each do |metric_labels, value|
+        str << "http_requests_total{"
+        str << "method=\"" << metric_labels.request_method << "\" "
+        str << "route=\"" << metric_labels.request_route << "\" "
+        str << "response_code=\"" << metric_labels.response_code << "\""
+        str << "} "
+        str << value << "\n"
+      end
+
+      Metrics::REQUEST_DURATION_SECONDS_SUMS.each do |metric_labels, value|
+        str << "http_request_duration_seconds_sum{"
+        str << "method=\"" << metric_labels.request_method << "\" "
+        str << "route=\"" << metric_labels.request_route << "\" "
+        str << "response_code=\"" << metric_labels.response_code << "\""
+        str << "} "
+        str << value << "\n"
+      end
+
+      Invidious::Jobs::StatisticsRefreshJob::STATISTICS_PROMETHEUS.each.each do |key, value|
+        str << key << " " << value << "\n"
+      end
+    end
   end
 
   # APIv1 currently uses the same logic for both
diff --git a/src/invidious/routing.cr b/src/invidious/routing.cr
index e7f0eee9..afd60f24 100644
--- a/src/invidious/routing.cr
+++ b/src/invidious/routing.cr
@@ -287,6 +287,7 @@ module Invidious::Routing
       # Misc
       get "/api/v1/stats", {{namespace}}::Misc, :stats
       if CONFIG.statistics_enabled
+        add_handler Metrics::METRICS_COLLECTOR
         get "/api/v1/metrics", {{namespace}}::Misc, :metrics
       end
       get "/api/v1/playlists/:plid", {{namespace}}::Misc, :get_playlist