electric

defmodule Electric.Telemetry.OpenTelemetry do
  @moduledoc """
  This module implements an API to cover parts of the code with tracing spans that are then
  exported using the OpenTelemetry protocol.

  [OpenTelemetry][1] is an observability framework that is widely supported by observability tools.

  This module's implementation is based on the [opentelemetry-erlang][2] suite of libraries.
  There is a rudimentary Elixir API there but it's incomplete and non-idiomatic. The idea with
  this module is to expose all of the functionality we're using in our code by wrapping
  opentelemetry-erlang's API.

  The configuration for OpenTelemetry export is located in `config/runtime.exs`.

  The API implemented here so far includes support for:

    - Defining a span to cover the execution of a piece of code. See `with_span/3`.

    - Propagating span context across Elixir processes, to allow for a span started in one
      process to be registered as a parent of a span started in a different process. See
      `get_current_context/1` and `set_current_context/1`.

    - Adding dynamic attributes to the current span, after it has already started. See
      `add_span_attributes/2`.

    - Recording an error or an exception as a span event. See `record_exception/4`.

  [1]: https://opentelemetry.io/docs/what-is-opentelemetry/
  [2]: https://github.com/open-telemetry/opentelemetry-erlang
  """

  require Logger
  require OpenTelemetry.SemanticConventions.Trace

  alias Electric.Telemetry.OptionalSpans

  @typep span_name :: String.t()
  @typep attr_name :: String.t()
  @typep span_attrs :: :opentelemetry.attributes_map()
  @typep span_ctx :: :opentelemetry.span_ctx()

  @doc """
  Create a span that starts at the current point in time and ends when `fun` returns.

  Returns the result of calling the function `fun`.

  Calling this function inside another span establishes a parent-child relationship between
  the two, as long as both calls happen within the same Elixir process. Use `get_current_context/1` for
  interprocess progragation of span context.

  The `stack_id` parameter must be set in root spans. For child spans the stack_id is optional
  and will be inherited from the parent span.
  """
  @spec with_span(span_name(), span_attrs(), String.t(), (-> t)) :: t when t: term
  def with_span(name, attributes, stack_id \\ nil, fun)
      when is_binary(name) and (is_list(attributes) or is_map(attributes)) do
    if OptionalSpans.include?(name) do
      do_with_span(name, attributes, stack_id, fun)
    else
      fun.()
    end
  end

  defp do_with_span(name, attributes, stack_id, fun) do
    stack_id = stack_id || get_from_baggage("stack_id")
    stack_attributes = get_stack_span_attrs(stack_id)
    all_attributes = stack_attributes |> Map.merge(Map.new(attributes))

    # This map is populated with default values that `:otel_tracer.with_span()` whould have set
    # anyway. But we're forced to do it here to avoid having like 50% of our code covered with
    # Dialyzer warnings (I dare you to try and only leave the `attributes` key here).
    span_opts = %{
      attributes: all_attributes,
      links: [],
      is_recording: true,
      start_time: :opentelemetry.timestamp(),
      kind: :internal
    }

    erlang_telemetry_event = [
      :electric | name |> String.split(".", trim: true) |> Enum.map(&String.to_atom/1)
    ]

    set_in_baggage("stack_id", stack_id)

    :telemetry.span(erlang_telemetry_event, all_attributes, fn ->
      fun_result = :otel_tracer.with_span(tracer(), name, span_opts, fn _span_ctx -> fun.() end)
      {fun_result, %{}}
    end)
  end

  @doc """
  Executes the provided function and records its duration in microseconds.
  The duration is added to the current span as a span attribute named with the given `name`.
  """
  @spec timed_fun(span_ctx() | nil, attr_name(), (-> term)) :: term
  def timed_fun(span_ctx \\ nil, name, fun) when is_binary(name) do
    {duration, result} = :timer.tc(fun)
    add_span_attributes(span_ctx, %{name => duration})
    result
  end

  @doc """
  Add dynamic attributes to the current span.

  For example, if a span is started prior to issuing a DB request, an attribute named
  `num_rows_fetched` can be added to it using this function once the DB query returns its
  result.
  """
  @spec add_span_attributes(span_ctx() | nil, span_attrs()) :: boolean()
  def add_span_attributes(span_ctx \\ nil, attributes) do
    span_ctx = span_ctx || current_span_context()
    :otel_span.set_attributes(span_ctx, attributes)
  end

  @doc """
  Store the telemetry span attributes in the persistent term for this stack.
  """
  @spec set_stack_span_attrs(String.t(), span_attrs()) :: :ok
  def set_stack_span_attrs(stack_id, attrs) do
    :persistent_term.put(:"electric_otel_attributes_#{stack_id}", Map.new(attrs))
  end

  @doc """
  Retrieve the telemetry span attributes from the persistent term for this stack.
  """
  @spec get_stack_span_attrs(String.t()) :: map()
  def get_stack_span_attrs(stack_id) do
    :persistent_term.get(:"electric_otel_attributes_#{stack_id}", %{})
  end

  @doc """
  Add an error event to the current span.
  """
  def record_exception(error_str, attributes \\ []) when is_binary(error_str) do
    add_exception_event("error", error_str, nil, attributes)
  end

  def record_exception(kind, error, stacktrace, attributes \\ []) when is_atom(kind) do
    type =
      if is_struct(error) do
        to_string(error.__struct__)
      else
        "error"
      end

    message = Exception.format(kind, error)
    add_exception_event(type, message, stacktrace, attributes)
  end

  defp add_exception_event(type, message, stacktrace, attributes) do
    semantic_attributes = [
      {OpenTelemetry.SemConv.ExceptionAttributes.exception_type(), type},
      {OpenTelemetry.SemConv.ExceptionAttributes.exception_message(), message},
      {OpenTelemetry.SemConv.ExceptionAttributes.exception_stacktrace(),
       Exception.format_stacktrace(stacktrace)}
    ]

    ctx = current_span_context()
    :otel_span.add_event(ctx, "exception", semantic_attributes ++ attributes)
    :otel_span.set_status(ctx, :error, message)
  end

  defp tracer, do: :opentelemetry.get_tracer()

  # Get the span and baggage context for the current process
  # Use this to pass the context to another process, see `set_current_context/1`
  def get_current_context do
    {current_span_context(), :otel_baggage.get_all()}
  end

  # Set the span and baggage context for the current process
  def set_current_context({span_ctx, baggage}) do
    :otel_tracer.set_current_span(span_ctx)
    :otel_baggage.set(baggage)
  end

  def set_in_baggage(key, value) do
    :otel_baggage.set(key, value)
  end

  def get_from_baggage(key) do
    case :otel_baggage.get_all() do
      %{^key => {value, _metadata}} -> value
      _ -> nil
    end
  end

  defp current_span_context do
    :otel_tracer.current_span_ctx()
  end
end