prometheus_exporter

#!/usr/bin/env ruby
# frozen_string_literal: true

require "bundler/setup"
require "prometheus_exporter"
require "prometheus_exporter/server"
require "json"
require "net/http"

require_relative "../lib/prometheus_exporter/custom_active_record_collector"

def fetch_ecs_task_id
  metadata_uri = ENV.fetch("ECS_CONTAINER_METADATA_URI_V4")
  max_retries = 12
  retry_delay = 10

  max_retries.times do |i|
    response = Net::HTTP.get(URI(metadata_uri))
    metadata = JSON.parse(response)
    task_arn = metadata.dig("Labels", "com.amazonaws.ecs.task-arn")
    return task_arn.split("/").last if task_arn
  rescue StandardError => e
    if i == max_retries - 1
      abort "Failed to fetch task details after #{max_retries} attempts: #{e.message}"
    end
    sleep retry_delay
  end
end

# An explicitly-set env var (including empty string) overrides the default.
# An empty value omits the label entirely — useful for the metrics task,
# whose series identity must stay stable across container replacements.
def resolve_label(env_key, &default)
  if ENV.key?(env_key)
    value = ENV[env_key]
    return value.empty? ? nil : value
  end

  default.call
end

labels = {
  "TaskId" => resolve_label("PROMETHEUS_TASK_ID_LABEL") do
    ARGV.include?("skip-server-labels") ? nil : fetch_ecs_task_id
  end,
  "ServiceName" => resolve_label("PROMETHEUS_SERVICE_NAME_LABEL") do
    ARGV.include?("skip-server-labels") ? nil : ENV["SERVICE_NAME"]
  end
}.compact

runner = PrometheusExporter::Server::Runner.new(
  port: 9394,
  bind: "0.0.0.0",
  label: labels,
  type_collectors: [PrometheusExporter::CustomActiveRecordCollector],
)

runner.start
sleep