Send logs from Elixir to Axiom

mix new axiom_elixir_demo --sup
cd axiom_elixir_demo

defmodule AxiomElixirDemo.MixProject do
  use Mix.Project

  def project do
    [
      app: :axiom_elixir_demo,
      version: "0.1.0",
      elixir: "~> 1.15",
      start_permanent: Mix.env() == :prod,
      deps: deps()
    ]
  end

  def application do
    [
      extra_applications: [:logger],
      mod: {AxiomElixirDemo.Application, []}
    ]
  end

  defp deps do
    [
      {:req, "~> 0.5"}
    ]
  end
end

mix deps.get

import Config

config :axiom_elixir_demo, AxiomLogger,
  token: System.fetch_env!("AXIOM_TOKEN"),
  dataset: System.fetch_env!("AXIOM_DATASET"),
  domain: System.fetch_env!("AXIOM_DOMAIN"),
  batch_size: 100,
  flush_interval: 2_000

export AXIOM_TOKEN=API_TOKEN
export AXIOM_DATASET=DATASET_NAME
export AXIOM_DOMAIN=AXIOM_DOMAIN

require Logger

Logger.debug("Checking details.")
Logger.info("User logged in.", user_id: "exampleUserId")
Logger.notice("User tried a feature.")
Logger.warning("Feature might not work as expected.")
Logger.error("Feature failed to load.", error_code: 500)
Logger.critical("Major issue with the app.")
Logger.alert("Immediate action needed.")
Logger.emergency("The app is down.")

00:30:05.795 [debug] Checking details.
00:30:05.800 [info] User logged in.
00:30:05.801 [notice] User tried a feature.
00:30:05.801 [warning] Feature might not work as expected.
00:30:05.801 [error] Feature failed to load.
00:30:05.801 [critical] Major issue with the app.
00:30:05.801 [alert] Immediate action needed.
00:30:05.801 [emergency] The app is down.

defmodule AxiomLogger.Handler do
  @moduledoc """
  An Erlang `:logger` handler that converts log events into JSON-friendly maps
  and hands them to `AxiomLogger.Buffer`, which batches them and sends them to Axiom.
  """

  # Default metadata keys that are noise in Axiom.
  @skip_metadata [:gl, :time, :report_cb, :erl_level, :initial_call, :process_label, :domain, :ancestors, :callers]

  # Called by :logger when the handler is added.
  def adding_handler(config), do: {:ok, config}

  # Called by :logger when the handler is removed. Flush so nothing is lost on shutdown.
  def removing_handler(_config), do: AxiomLogger.Buffer.flush()

  # Never forward the buffer's own log events. This prevents infinite loops.
  def log(%{meta: %{axiom_internal: true}}, _config), do: :ok

  # Called by :logger in the process that emitted the log event, so keep it cheap:
  # build the event and hand it to the buffer.
  def log(%{level: level, msg: msg, meta: meta}, _config) do
    AxiomLogger.Buffer.push(build_event(level, msg, meta))
  end

  defp build_event(level, msg, meta) do
    timestamp =
      meta
      |> Map.get(:time, System.system_time(:microsecond))
      |> DateTime.from_unix!(:microsecond)
      |> DateTime.to_iso8601()

    %{"_time" => timestamp, "level" => Atom.to_string(level), "metadata" => build_metadata(meta)}
    |> put_message(msg)
  end

  # Plain string messages, for example Logger.info("User logged in.")
  defp put_message(event, {:string, chardata}) do
    Map.put(event, "message", IO.chardata_to_string(chardata))
  end

  # Structured reports, for example Logger.info(%{event: "order_placed", order_id: 42})
  defp put_message(event, {:report, report}) when is_map(report) do
    Map.put(event, "report", sanitize(report))
  end

  defp put_message(event, {:report, report}) when is_list(report) do
    if Keyword.keyword?(report) do
      Map.put(event, "report", report |> Map.new() |> sanitize())
    else
      Map.put(event, "message", inspect(report))
    end
  end

  # Erlang-style format strings, for example :logger.info(~c"~p items", [3])
  defp put_message(event, {format, args}) do
    Map.put(event, "message", format |> :io_lib.format(args) |> IO.chardata_to_string())
  rescue
    _ -> Map.put(event, "message", inspect({format, args}))
  end

  defp build_metadata(meta) do
    {mfa, meta} = Map.pop(meta, :mfa)
    {file, meta} = Map.pop(meta, :file)

    meta
    |> Map.drop(@skip_metadata)
    |> sanitize()
    |> maybe_put("file", file && to_string(file))
    |> put_mfa(mfa)
  end

  defp put_mfa(metadata, {module, function, arity}) do
    Map.merge(metadata, %{"module" => inspect(module), "function" => "#{function}/#{arity}"})
  end

  defp put_mfa(metadata, _), do: metadata

  defp maybe_put(map, _key, nil), do: map
  defp maybe_put(map, key, value), do: Map.put(map, key, value)

  # Convert any Elixir term into something the JSON encoder accepts.
  defp sanitize(value) when is_binary(value) do
    if String.valid?(value), do: value, else: inspect(value)
  end

  defp sanitize(value) when is_number(value) or is_boolean(value) or is_nil(value), do: value
  defp sanitize(value) when is_atom(value), do: Atom.to_string(value)
  defp sanitize(%DateTime{} = value), do: DateTime.to_iso8601(value)
  defp sanitize(%NaiveDateTime{} = value), do: NaiveDateTime.to_iso8601(value)
  defp sanitize(value) when is_struct(value), do: inspect(value)

  defp sanitize(value) when is_map(value) do
    Map.new(value, fn {key, val} -> {sanitize_key(key), sanitize(val)} end)
  end

  defp sanitize(value) when is_list(value) do
    if value != [] and Keyword.keyword?(value) do
      value |> Map.new() |> sanitize()
    else
      Enum.map(value, &sanitize/1)
    end
  end

  # PIDs, references, tuples, functions, and ports
  defp sanitize(value), do: inspect(value)

  defp sanitize_key(key) when is_atom(key), do: Atom.to_string(key)
  defp sanitize_key(key) when is_binary(key), do: key
  defp sanitize_key(key), do: inspect(key)
end

defmodule AxiomLogger.Buffer do
  @moduledoc """
  Collects log events and sends them to the Axiom ingest API in batches.

  Events are flushed when the batch reaches `:batch_size` events or when
  `:flush_interval` milliseconds have passed, whichever comes first.
  """
  use GenServer
  require Logger

  @default_batch_size 100
  @default_flush_interval 2_000

  def start_link(opts) do
    GenServer.start_link(__MODULE__, opts, name: __MODULE__)
  end

  @doc "Adds an event to the buffer. Never blocks the caller."
  def push(event), do: GenServer.cast(__MODULE__, {:push, event})

  @doc "Sends all buffered events immediately and waits for the request to finish."
  def flush, do: GenServer.call(__MODULE__, :flush, 30_000)

  @impl true
  def init(opts) do
    # Run terminate/2 on shutdown so the last batch is sent.
    Process.flag(:trap_exit, true)

    # Mark this process so AxiomLogger.Handler ignores anything it logs.
    Logger.metadata(axiom_internal: true)

    domain = Keyword.fetch!(opts, :domain)
    dataset = Keyword.fetch!(opts, :dataset)

    state = %{
      url: "https://#{domain}/v1/ingest/#{URI.encode(dataset)}",
      token: Keyword.fetch!(opts, :token),
      batch_size: Keyword.get(opts, :batch_size, @default_batch_size),
      flush_interval: Keyword.get(opts, :flush_interval, @default_flush_interval),
      events: [],
      count: 0
    }

    {:ok, schedule_flush(state)}
  end

  @impl true
  def handle_cast({:push, event}, state) do
    state = %{state | events: [event | state.events], count: state.count + 1}

    if state.count >= state.batch_size do
      {:noreply, do_flush(state)}
    else
      {:noreply, state}
    end
  end

  @impl true
  def handle_call(:flush, _from, state), do: {:reply, :ok, do_flush(state)}

  @impl true
  def handle_info(:flush, state), do: {:noreply, state |> do_flush() |> schedule_flush()}

  @impl true
  def terminate(_reason, state) do
    do_flush(state)
    :ok
  end

  defp schedule_flush(state) do
    Process.send_after(self(), :flush, state.flush_interval)
    state
  end

  defp do_flush(%{events: []} = state), do: state

  defp do_flush(state) do
    events = Enum.reverse(state.events)

    # retry: :transient retries HTTP 408/429/5xx responses and connection errors.
    case Req.post(state.url, json: events, auth: {:bearer, state.token}, retry: :transient) do
      {:ok, %Req.Response{status: 200, body: %{"failed" => 0}}} ->
        :ok

      {:ok, %Req.Response{status: 200, body: body}} ->
        Logger.warning("Axiom rejected some events: #{inspect(body["failures"])}")

      {:ok, %Req.Response{status: status, body: body}} ->
        Logger.error("Axiom ingest failed with HTTP #{status}: #{inspect(body)}")

      {:error, error} ->
        Logger.error("Axiom ingest request failed: #{Exception.message(error)}")
    end

    %{state | events: [], count: 0}
  end
end

defmodule AxiomElixirDemo.Application do
  @moduledoc false
  use Application

  @impl true
  def start(_type, _args) do
    children = [
      {AxiomLogger.Buffer, Application.fetch_env!(:axiom_elixir_demo, AxiomLogger)}
    ]

    opts = [strategy: :one_for_one, name: AxiomElixirDemo.Supervisor]
    {:ok, pid} = Supervisor.start_link(children, opts)

    # Attach the handler after the buffer process is running.
    :ok = :logger.add_handler(:axiom, AxiomLogger.Handler, %{level: :debug})

    {:ok, pid}
  end

  @impl true
  def prep_stop(state) do
    # Removing the handler flushes buffered events before the supervision tree stops.
    :logger.remove_handler(:axiom)
    state
  end
end

defmodule AxiomElixirDemo do
  @moduledoc "Emits a log event at every level so you can see them arrive in Axiom."
  require Logger

  def run do
    Logger.debug("Checking details.", action: "detail_check", status: "initiated")
    Logger.info("User logged in.", user_id: "exampleUserId", method: "standard_login")
    Logger.notice("User tried a feature.", feature: "experimental_feature_x", status: "trial")
    Logger.warning("Feature might not work as expected.", feature: "experimental_feature", stage: "beta")
    Logger.error("Feature failed to load.", feature: "feature_y", error_code: 500)
    Logger.critical("Major issue with the app.", system: "payment_processing", error: "service_unavailable")
    Logger.alert("Immediate action needed.", issue: "security", severity: "high")
    Logger.emergency("The app is down.", system: "entire_application", status: "offline")

    # Structured report: the map becomes the `report` field in Axiom.
    Logger.info(%{event: "order_placed", order_id: 1234, total: 59.99, currency: "USD"})

    # Send whatever is still buffered before the script exits.
    AxiomLogger.Buffer.flush()
  end
end

mix run -e "AxiomElixirDemo.run()"

['DATASET_NAME']
| where level in ("error", "critical", "alert", "emergency")
| project _time, level, message, ['metadata.module'], ['metadata.line']
| order by _time desc