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
  """

  alias Electric.Telemetry.Sampler
  alias Electric.Telemetry.IntervalTimer

  require Record

  Record.defrecordp(
    :span_ctx_record,
    :span_ctx,
    Record.extract(:span_ctx, from_lib: "opentelemetry_api/include/opentelemetry.hrl")
  )

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

  @typedoc """
  Span + baggage pair returned by `get_current_context/0` and consumed by `set_current_context/1`.
  """
  @type otel_ctx :: {span_ctx() | :undefined, :otel_baggage.t()}

  @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() | nil, (-> 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
    do_with_span(name, attributes, stack_id, fun, Sampler.include_span?(name))
  end

  @doc """
  Creates a span providing there is a parent span in the current context.
  If there is no parent span, the function `fun` is called without creating a span.

  This is necessary for the custom way we do sampling, if the parent span is not sampled, the child span
  will not be created either.
  """
  def with_child_span(name, attributes, stack_id \\ nil, fun) do
    do_with_span(
      name,
      attributes,
      stack_id,
      fun,
      in_span_context?() && Sampler.include_span?(name)
    )
  end

  defp do_with_span(name, attributes, stack_id, fun, include_otel_span?) do
    erlang_telemetry_event = [
      :electric | name |> String.split(".", trim: true) |> Enum.map(&String.to_atom/1)
    ]

    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))

    :telemetry.span(erlang_telemetry_event, all_attributes, fn ->
      fun_result =
        if include_otel_span? do
          with_otel_span(name, all_attributes, stack_id, fun)
        else
          fun.()
        end

      {fun_result, %{}}
    end)
  end

  defp with_otel_span(name, attributes, stack_id, fun) do
    # 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: attributes,
      links: [],
      is_recording: true,
      start_time: :opentelemetry.timestamp(),
      kind: :internal
    }

    set_in_baggage("stack_id", stack_id)

    :otel_tracer.with_span(tracer(), name, span_opts, fn _span_ctx -> fun.() end)
  end

  @doc """
  A thin wrapper around `:telemetry.execute/3` that adds the span attributes for the current
  stack to the metadata.
  """
  @spec execute(
          :telemetry.event_name(),
          :telemetry.event_measurements() | :telemetry.event_value(),
          :telemetry.event_metadata()
        ) :: :ok
  def execute(event_name, measurements, metadata) do
    stack_id = metadata[:stack_id] || get_from_baggage("stack_id")
    metadata = Map.merge(metadata, get_stack_span_attrs(stack_id))
    :telemetry.execute(event_name, measurements, metadata)
  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 """
  Execute `fun`, recording its duration and process-memory footprint as
  attributes on the current span.

  This is a flattened alternative to `with_span/4` for code that is strictly
  1:1 with its enclosing span: instead of emitting a separate child span, the
  same information is recorded as attributes on the current span, prefixed
  with `name`:

    * `<name>.duration_ms` — execution time of `fun` in milliseconds
    * `<name>.memory.start.*` / `<name>.memory.end.*` — the process-memory
      footprint attributes from `process_memory_attributes/1`, captured before
      and after `fun` runs

  The attributes are recorded even when `fun` raises, mirroring how a child
  span would still have been ended with its duration. The exception itself
  propagates up unchanged, to be recorded on the enclosing span by its error
  handling (see `record_exception/4`).

  When there is no active span context this is a no-op apart from calling
  `fun`, avoiding the `Process.info/2` cost on unsampled requests.
  """
  @spec with_flattened_span(attr_name(), (-> t)) :: t when t: term
  def with_flattened_span(name, fun) when is_binary(name) do
    if in_span_context?() do
      start_attrs = prefix_attrs(name, process_memory_attributes(:start))
      start_time = System.monotonic_time()

      try do
        fun.()
      after
        duration_ms =
          System.convert_time_unit(System.monotonic_time() - start_time, :native, :microsecond) /
            1000

        end_attrs = prefix_attrs(name, process_memory_attributes(:end))

        start_attrs
        |> Map.merge(end_attrs)
        |> Map.put("#{name}.duration_ms", duration_ms)
        |> add_span_attributes()
      end
    else
      fun.()
    end
  end

  defp prefix_attrs(prefix, attrs) do
    Map.new(attrs, fn {key, value} -> {"#{prefix}.#{key}", value} end)
  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 """
  Build a map of current-process memory-footprint attributes suitable for use
  as span attributes.

  Captures two values via `Process.info/2`:
    * `process_bytes` — total memory occupied by the process (heap, stack,
      message queue, GC overhead, etc.)
    * `binary_bytes` — sum of the sizes of refc binaries referenced by the
      process; this is what tends to dominate memory in shape requests that
      buffer large response bodies.

  The `phase` argument determines the attribute key prefix: `:start` produces
  `memory.start.process_bytes` and `memory.start.binary_bytes`, `:end` produces
  `memory.end.process_bytes` and `memory.end.binary_bytes`.

  Caveat: `Process.info(self(), :binary)` lists one entry per reference, so a
  refc binary that the process references multiple times is counted multiple
  times in `binary_bytes`. This is the conventional approximation (Recon's
  `recon:bin_leak/1` does the same) and is fine for tracking growth across a
  span — but it should not be used as a tight per-process memory budget.
  `:erlang.memory(:binary)` gives a deduplicated VM-wide figure if you need
  one.
  """
  @spec process_memory_attributes(:start | :end) :: %{String.t() => non_neg_integer()}
  def process_memory_attributes(phase) when phase in [:start, :end] do
    [memory: process_bytes, binary: binaries] = Process.info(self(), [:memory, :binary])
    binary_bytes = Enum.reduce(binaries, 0, fn {_ref, size, _count}, acc -> acc + size end)

    case phase do
      :start ->
        %{
          "memory.start.process_bytes" => process_bytes,
          "memory.start.binary_bytes" => binary_bytes
        }

      :end ->
        %{
          "memory.end.process_bytes" => process_bytes,
          "memory.end.binary_bytes" => binary_bytes
        }
    end
  end

  @doc """
  Add current-process memory-footprint attributes (see
  `process_memory_attributes/1`) to the current span.

  No-op when called outside a span context, avoiding the `Process.info/2` cost
  on unsampled requests.
  """
  @spec add_process_memory_attributes(:start | :end) :: boolean()
  def add_process_memory_attributes(phase) when phase in [:start, :end] do
    if in_span_context?() do
      add_span_attributes(process_memory_attributes(phase))
    else
      false
    end
  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 """
  Records that an interval with the given `interval_name` has started.

  This is useful if you want to find out which part of a process took
  the longest time. It works out simpler than wrapping each part of
  the process in a timer, and guarentees no gaps in the timings.

  Once a number of intervals have been started, call
  `stop_and_save_intervals()` to record the interval timings as
  attributes in the current span.

  e.g.

  ```elixir
  OpenTelemetry.start_interval(:quick_sleep.duration_µs)
  Process.sleep(1)
  OpenTelemetry.start_interval(:longer_sleep.duration_µs)
  Process.sleep(2)
  OpenTelemetry.stop_and_save_intervals(total_attribute: "total_sleep_µs")
  ```
  will add the following attributes to the current span:
    quick_sleep.duration_µs: 1000
    longer_sleep.duration_µs: 2000
    total_sleep_µs: 3000
  """
  @spec start_interval(atom()) :: :ok
  def start_interval(interval_name) when is_atom(interval_name) do
    IntervalTimer.start_interval(get_interval_timer(), interval_name)
    |> set_interval_timer()

    :ok
  end

  @doc """
  Records the interval timings as attributes in the current span
  and wipes the interval timer from process memory.

  Options:
  - `:timer` - the interval timer to use. If not provided, the timer
    is extracted from the process memory.
  - `:total_attribute` - the name of the attribute to store the total
    time across all intervals. If not provided no total time is recorded.
  """
  def stop_and_save_intervals(opts) do
    timer = opts[:timer] || extract_interval_timer()
    durations = IntervalTimer.durations(timer)

    total_attribute =
      case opts[:total_attribute] do
        nil -> []
        attr_name -> [{attr_name, IntervalTimer.total_time(durations)}]
      end

    add_span_attributes(
      total_attribute ++
        for {interval_name, duration} <- durations do
          {interval_name, duration}
        end
    )
  end

  @interval_timer_key :electric_otel_interval_timer

  @doc """
  Set the interval timer for the current process.
  """
  @spec set_interval_timer(IntervalTimer.t()) :: :ok
  def set_interval_timer(timer) do
    Process.put(@interval_timer_key, timer)
  end

  @doc """
  Wipe the current interval timer from process memory.
  """
  def wipe_interval_timer do
    Process.delete(@interval_timer_key)
  end

  @doc """
  Removes the current interval timer from prcess memory and returns it.

  Useful if you want to time intervals over multiple processes,
  extract the timer, pass it to another process, and then
  use `set_interval_timer/1` to restore it in the new process.
  """
  @spec extract_interval_timer() :: IntervalTimer.t()
  def extract_interval_timer do
    timer = get_interval_timer()
    wipe_interval_timer()
    timer
  end

  defp get_interval_timer do
    Process.get(@interval_timer_key, [])
  end

  @doc """
  Look up the value of the member with the given key in the W3C tracestate carried
  by the given span context.

  Returns the member's value as a string, or `nil` when the tracestate has no such
  member.
  """
  @spec tracestate_value(span_ctx(), String.t()) :: String.t() | nil
  def tracestate_value(span_ctx, key) when Record.is_record(span_ctx, :span_ctx) do
    case :otel_tracestate.get(key, span_ctx_record(span_ctx, :tracestate)) do
      value when is_binary(value) and value != "" -> value
      [_ | _] = value -> IO.iodata_to_binary(value)
      _ -> nil
    end
  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

  defp in_span_context? do
    current_span_context() != :undefined
  end
end