kubernetes_probes

defmodule KubernetesProbes.Drainer do
  @moduledoc """
  Holds the k8s readiness-probe status and enforces a drain window on shutdown.

  On `SIGTERM`, OTP walks the supervision tree in reverse start order and
  terminates children in reverse order. This `GenServer` **must be the last
  child** of your application supervisor so it terminates first. Its
  `terminate/2` flips the `/probe/readiness` probe to 503 via `:persistent_term`
  and sleeps for the drain window, giving Kubernetes time to stop routing
  traffic and allowing in-flight requests to finish before other resources
  (Repo, Endpoint, etc.) are torn down.

  ## Usage

      # lib/my_app/application.ex — last entry in children
      children = [
        MyApp.Repo,
        MyAppWeb.Endpoint,
        {KubernetesProbes.Drainer, otp_app: :my_app}
      ]

  ## Configuration

      # config/config.exs — production default
      config :my_app, KubernetesProbes.Drainer, wait: 20_000

      # config/dev.exs — avoid 20s hang on Ctrl-C
      config :my_app, KubernetesProbes.Drainer, wait: 100

      # config/test.exs — avoid slow suite teardown
      config :my_app, KubernetesProbes.Drainer, wait: 10

  ## Options

    * `:otp_app` — the OTP application to read configuration from (required).
    * `:wait` — drain window in milliseconds. Can be set via config (see above)
      or passed directly to override config: `{KubernetesProbes.Drainer, otp_app: :my_app, wait: 5_000}`.
      Defaults to `20_000`.
  """

  use GenServer
  require Logger

  @default_wait 20_000
  @shutdown_slack 5_000

  def child_spec(opts) do
    wait = resolve_wait(opts)

    %{
      id: __MODULE__,
      start: {__MODULE__, :start_link, [opts]},
      shutdown: wait + @shutdown_slack,
      restart: :permanent,
      type: :worker
    }
  end

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

  @spec status() :: :running | :stopping
  def status, do: :persistent_term.get({__MODULE__, :status}, :running)

  @impl true
  def init(opts) do
    Process.flag(:trap_exit, true)
    :persistent_term.put({__MODULE__, :status}, :running)
    {:ok, %{wait: resolve_wait(opts)}}
  end

  @impl true
  def terminate(_reason, %{wait: wait}) do
    Logger.info("#{__MODULE__}: draining for #{wait}ms")
    :persistent_term.put({__MODULE__, :status}, :stopping)
    Process.sleep(wait)
    :ok
  end

  defp resolve_wait(opts) do
    otp_app = Keyword.get(opts, :otp_app)
    app_wait = otp_app && Application.get_env(otp_app, __MODULE__, []) |> Keyword.get(:wait)
    Keyword.get(opts, :wait, app_wait || @default_wait)
  end
end