kubernetes_probes

defmodule KubernetesProbes.Plug do
  @moduledoc """
  Handles k8s liveness and readiness probes. Use as the first plug in `endpoint.ex`.

  - `GET /probe/liveness` → 200 while the BEAM is running.
  - `GET /probe/readiness` → 200 when `KubernetesProbes.Drainer` is `:running` and the
    optional `ready?` function returns `true`; 503 otherwise.

  ## Usage

      # lib/my_app_web/endpoint.ex — before all other plugs
      plug KubernetesProbes.Plug

      # With a custom readiness check (e.g. database connectivity)
      plug KubernetesProbes.Plug, ready?: &MyApp.repos_ready?/0

      # With custom probe paths
      plug KubernetesProbes.Plug, liveness_path: "/healthz", readiness_path: "/readyz"

  ## Options

    * `:ready?` — zero-arity boolean function called on each readiness request
      while the drainer is `:running`. Must be a remote function capture
      (`&MyApp.repos_ready?/0`), not an anonymous function. Defaults to
      `&KubernetesProbes.Plug.always_ready/0` (always returns `true`).
    * `:liveness_path` — path for the liveness probe. Defaults to `"/probe/liveness"`.
    * `:readiness_path` — path for the readiness probe. Defaults to `"/probe/readiness"`.
  """

  import Plug.Conn

  @behaviour Plug

  def always_ready, do: true

  @impl true
  def init(opts) do
    ready? = Keyword.get(opts, :ready?, &__MODULE__.always_ready/0)
    liveness_path = Keyword.get(opts, :liveness_path, "/probe/liveness")
    readiness_path = Keyword.get(opts, :readiness_path, "/probe/readiness")

    unless is_function(ready?, 0) do
      raise ArgumentError, ":ready? must be a zero-arity function"
    end

    unless is_binary(liveness_path) do
      raise ArgumentError, ":liveness_path must be a string"
    end

    unless is_binary(readiness_path) do
      raise ArgumentError, ":readiness_path must be a string"
    end

    %{ready?: ready?, liveness_path: liveness_path, readiness_path: readiness_path}
  end

  @impl true
  def call(%Plug.Conn{request_path: path} = conn, %{liveness_path: path}) do
    conn |> send_resp(200, "") |> halt()
  end

  def call(%Plug.Conn{request_path: path} = conn, %{readiness_path: path, ready?: ready?}) do
    case KubernetesProbes.Drainer.status() do
      :stopping -> conn |> send_resp(503, "Draining") |> halt()
      :running -> conn |> respond_by_readiness(ready?) |> halt()
    end
  end

  def call(conn, _opts), do: conn

  defp respond_by_readiness(conn, ready?) do
    case ready?.() do
      true -> send_resp(conn, 200, "Serving")
      false -> send_resp(conn, 503, "Not ready")
    end
  end
end