hexagon_tpu

defmodule HexagonTpu.Stats do
  @moduledoc """
  Observability for the NIF layer: native resource accounting (leak guard),
  dirty-scheduler utilization, and memory snapshots.

  The three signals to watch in production:

    * `native/0` — `runtimes_alive`/`contexts_alive`/`graphs_alive` must stay
      bounded; monotonic growth means QNN resources are being leaked
      (references held forever, or destructors never running).
    * `dirty_utilization/1` — inference runs on dirty schedulers; sustained
      saturation there starves other NIFs and file I/O.
    * `memory/0` — output tensors are refc binaries; growth in `:binary`
      without matching workload growth points at tensors being retained.

  `HexagonTpu.Monitor` samples all of these periodically and emits telemetry.
  """

  alias HexagonTpu.Nif

  @doc """
  Native counters from the NIF:

    * `runtimes_alive` / `contexts_alive` / `graphs_alive` — QNN resources not
      yet destructed (leak guard: should return to baseline after GC)
    * `*_created` — lifetime totals
    * `executes` / `execute_errors` — inference counters
    * `execute_ns_total` / `execute_ns_last` — pure QNN `graphExecute` time
  """
  def native, do: Nif.stats()

  @doc """
  Samples dirty CPU / dirty IO scheduler utilization over `sample_ms`
  (default 1000) using microstate accounting.

  Returns `%{dirty_cpu: float, dirty_io: float, scheduler: float}` where each
  value is 0.0..1.0 busy share for that scheduler class.

  Microstate accounting is a VM-global flag: concurrent samplers (two
  Monitors, an operator using `msacc` in a remote shell) corrupt each
  other's windows. Run one sampler per node — `HexagonTpu.Monitor` is the
  intended one.
  """
  def dirty_utilization(sample_ms \\ 1_000) do
    :erlang.system_flag(:microstate_accounting, :reset)
    :erlang.system_flag(:microstate_accounting, true)
    Process.sleep(sample_ms)
    :erlang.system_flag(:microstate_accounting, false)

    by_type =
      :erlang.statistics(:microstate_accounting)
      |> Enum.group_by(& &1.type)
      |> Map.new(fn {type, threads} -> {type, busy_share(threads)} end)

    %{
      dirty_cpu: Map.get(by_type, :dirty_cpu_scheduler, 0.0),
      dirty_io: Map.get(by_type, :dirty_io_scheduler, 0.0),
      scheduler: Map.get(by_type, :scheduler, 0.0)
    }
  end

  @doc """
  Memory snapshot relevant to the NIF: refc binaries (output tensors live
  here) and total BEAM memory.
  """
  def memory do
    %{
      binary: :erlang.memory(:binary),
      total: :erlang.memory(:total),
      system: :erlang.memory(:system)
    }
  end

  @doc """
  Runs `fun`, forces a full GC sweep, and returns
  `{result, before_stats, after_stats}` of the native counters. Useful in
  tests to assert no resources leaked across an operation.
  """
  def with_leak_check(fun) do
    before_stats = native()
    result = fun.()
    sweep()
    {result, before_stats, native()}
  end

  @doc """
  Forces a GC of all processes so NIF resource destructors for dropped
  references run.
  """
  def sweep do
    Enum.each(Process.list(), &:erlang.garbage_collect/1)
    # Destructors run asynchronously on resource release; give them a beat.
    Process.sleep(50)
    :ok
  end

  defp busy_share(threads) do
    {busy, total} =
      Enum.reduce(threads, {0, 0}, fn %{counters: counters}, {busy, total} ->
        sleep = Map.get(counters, :sleep, 0)
        sum = counters |> Map.values() |> Enum.sum()
        {busy + (sum - sleep), total + sum}
      end)

    if total == 0, do: 0.0, else: busy / total
  end
end