circuitbreaker.ex

defmodule CircuitBreaker do
  @moduledoc """
  `CircuitBreaker` is a `decorator` module for supervisor children that changes the default supervisor restart policy.

  Instead of instantly restarting an exited children, like a normal supervisor
  would do, `CircuitBreaker` will instead change its status to `in_break` (i.e.
  the worker is `unnavailable`) and schedule a restart of the children after a
  configured `delay`. If the children fails to start, the status is kept
  `in_break` and a new restart is scheduled. If the children starts
  successfully, then the status is changed back to `normal`.

  Clients can check whether a worker is `in_break?/1` (unavailable) and also
  force the worker to go into `in_break` via `break!/1`. 
  """

  use Supervisor

  def start_link(child_spec, breaker_ref \\ nil, opts \\ []) do
    breaker_ref = breaker_ref || child_spec

    init_args = %{breaker_ref: breaker_ref, child_spec: child_spec}
    opts = Keyword.put_new(opts, :name, breaker_ref)

    Supervisor.start_link(__MODULE__, init_args, opts)
  end

  @impl true
  def init(_args = %{breaker_ref: breaker_ref, child_spec: child_spec}) when breaker_ref != nil do
    :ets.new(breaker_ref, [:set, :public, :named_table])

    children = [
      Supervisor.child_spec(child_spec, id: :worker, restart: :transient),
      {CircuitBreaker.Restarter, [breaker_ref: breaker_ref]}
    ]

    Supervisor.init(children, strategy: :one_for_all)
  end

  @spec break!(atom()) :: :ok
  def break!(breaker_ref) do
    :ok = set_break(breaker_ref, true)
  end

  @spec restart_worker(atom()) :: any()
  def restart_worker(breaker_ref) do
    case Supervisor.restart_child(breaker_ref, :worker) do
      {:ok, _pid} ->
        :ok = set_break(breaker_ref, false)

      {:ok, _pid, _info} ->
        :ok = set_break(breaker_ref, false)

      otherwise ->
        otherwise
    end
  end

  @spec in_break?(atom()) :: boolean()
  def in_break?(breaker_ref) do
    case :ets.lookup(breaker_ref, :in_break?) do
      [{:in_break?, in_break?}] -> in_break?
      [] -> false
    end
  end

  @spec set_break(atom(), boolean()) :: :ok
  defp set_break(breaker_ref, value) do
    true = :ets.insert(breaker_ref, {:in_break?, value})
    :ok
  end
end

defmodule CircuitBreaker.Restarter do
  use GenServer, restart: :permanent, id: :restarter

  def start_link(args, opts \\ []),
    do: GenServer.start_link(__MODULE__, args, opts)

  @impl true
  def init(args) do
    Process.flag(:trap_exit, true)
    breaker_ref = args[:breaker_ref] || raise ArgumentError

    if CircuitBreaker.in_break?(breaker_ref) do
      schedule_worker_restart()
    end

    {:ok, %{breaker_ref: breaker_ref}}
  end

  @impl true
  def terminate(_reason, _state = %{breaker_ref: breaker_ref}) do
    :ok = CircuitBreaker.break!(breaker_ref)
  end

  @impl true
  def handle_info(:schedule_worker_restart, state = %{breaker_ref: breaker_ref}) do
    unless CircuitBreaker.restart_worker(breaker_ref) == :ok do
      schedule_worker_restart()
    end
    {:noreply, state}
  end

  defp schedule_worker_restart do
    Process.send_after(self(), :schedule_worker_restart, 1_000)
  end
end