From e165dd0fdb24c2de678ab59706c96d483e63add7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Niemier?= <lukasz@niemier.pl>
Date: Fri, 12 Jul 2024 09:16:05 +0200
Subject: [PATCH] ft: implement pluggable storage

Preliminary work for supporting different storages beyond
`telemetry_metrics_prometheus_core`.
---
 lib/prom_ex.ex              | 23 ++++++++++-------------
 lib/prom_ex/storage.ex      | 20 ++++++++++++++++++++
 lib/prom_ex/storage/core.ex | 25 +++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 13 deletions(-)
 create mode 100644 lib/prom_ex/storage.ex
 create mode 100644 lib/prom_ex/storage/core.ex

diff --git a/lib/prom_ex.ex b/lib/prom_ex.ex
index bca4973..dd9b683 100644
--- a/lib/prom_ex.ex
+++ b/lib/prom_ex.ex
@@ -106,8 +106,6 @@ defmodule PromEx do
 
   alias Telemetry.Metrics.{Counter, Distribution, LastValue, Sum, Summary}
 
-  alias TelemetryMetricsPrometheus.Core
-
   @type telemetry_metrics() :: Counter.t() | Distribution.t() | LastValue.t() | Sum.t() | Summary.t()
   @type measurements_mfa() :: {module(), atom(), list()}
 
@@ -122,10 +120,9 @@ defmodule PromEx do
   @spec get_metrics(prom_ex_module :: module()) :: String.t() | :prom_ex_down
   def get_metrics(prom_ex_module) do
     prom_ex_process_name = prom_ex_module.__metrics_collector_name__()
+    store = prom_ex_module.__store__()
 
-    if Process.whereis(prom_ex_process_name),
-      do: Core.scrape(prom_ex_process_name),
-      else: :prom_ex_down
+    PromEx.Storage.scrape(store, prom_ex_process_name)
   end
 
   @callback init_opts :: PromEx.Config.t()
@@ -147,6 +144,8 @@ defmodule PromEx do
           raise "Failed to initialize #{inspect(calling_module)} due to missing :otp_app option"
       end
 
+    store = Keyword.get(opts, :store, PromEx.Storage.Core)
+
     # Generate process names under calling module namespace
     ets_cron_flusher_name = Module.concat([calling_module, ETSCronFlusher])
     manual_metrics_name = Module.concat([calling_module, ManualMetricsManager])
@@ -202,7 +201,7 @@ defmodule PromEx do
           children =
             []
             |> PromEx.ets_cron_flusher_child_spec(__MODULE__, ets_flush_interval, unquote(ets_cron_flusher_name))
-            |> PromEx.metrics_collector_child_spec(telemetry_metrics, unquote(metrics_collector_name))
+            |> PromEx.metrics_collector_child_spec(unquote(store), telemetry_metrics, unquote(metrics_collector_name))
             |> PromEx.manual_metrics_child_spec(
               manual_metrics,
               manual_metrics_start_delay,
@@ -334,6 +333,9 @@ defmodule PromEx do
       @doc false
       def __ets_cron_flusher_name__, do: unquote(ets_cron_flusher_name)
 
+      @doc false
+      def __store__, do: unquote(store)
+
       defoverridable PromEx
     end
   end
@@ -372,13 +374,8 @@ defmodule PromEx do
   end
 
   @doc false
-  def metrics_collector_child_spec(acc, metrics, process_name) do
-    spec = {
-      Core,
-      name: process_name, metrics: metrics, require_seconds: false, consistent_units: true, start_async: false
-    }
-
-    [spec | acc]
+  def metrics_collector_child_spec(acc, store, metrics, process_name) do
+    [PromEx.Storage.child_spec(store, process_name, metrics) | acc]
   end
 
   @doc false
diff --git a/lib/prom_ex/storage.ex b/lib/prom_ex/storage.ex
new file mode 100644
index 0000000..609fb5f
--- /dev/null
+++ b/lib/prom_ex/storage.ex
@@ -0,0 +1,20 @@
+defmodule PromEx.Storage do
+  @moduledoc """
+  Storage definition behaviour.
+  """
+
+  @doc """
+  Gather metrics for given collector with given `name`.
+  """
+  @callback scrape(name :: atom()) :: iodata() | :prom_ex_down
+
+  @doc """
+  Define child specs for gatherer process.
+  """
+  @callback child_spec(atom(), Telemetry.Metrics.metrics()) :: Supervisor.child_spec()
+
+  def scrape(mod, name), do: mod.scrape(name)
+
+  def child_spec(mod, name, metrics),
+    do: mod.child_spec(name, metrics)
+end
diff --git a/lib/prom_ex/storage/core.ex b/lib/prom_ex/storage/core.ex
new file mode 100644
index 0000000..d6df3c9
--- /dev/null
+++ b/lib/prom_ex/storage/core.ex
@@ -0,0 +1,25 @@
+defmodule PromEx.Storage.Core do
+  @behaviour PromEx.Storage
+
+  alias TelemetryMetricsPrometheus.Core
+
+  @impl true
+  def scrape(name) do
+    if Process.whereis(name),
+      do: Core.scrape(name),
+      else: :prom_ex_down
+  end
+
+  @impl true
+  def child_spec(name, metrics) do
+    opts = [
+      name: name,
+      metrics: metrics,
+      require_seconds: false,
+      consistent_units: true,
+      start_async: false
+    ]
+
+    Core.child_spec(opts)
+  end
+end