EventStore.AdvisoryLocks hangs for good when its connection pool dies while new connections are refused

# After the advisory lock connection pool of an event store dies while PostgreSQL refuses new
# connections, EventStore.AdvisoryLocks can end up waiting for a connection checkout that nothing
# will answer. Every subscription of the event store crashes on its 15 s lock call, and no
# subscription can subscribe until AdvisoryLocks or the event store restarts.
#
# Every connection of the event store goes through a small TCP proxy in this script. The proxy
# either forwards new connections or refuses them; connections it already relays keep working.
# Two persistent subscriptions take their advisory locks. The script then makes the proxy refuse
# new connections, terminates the backend that holds the advisory locks, forwards again 45 s later
# and prints the state 60 s after that. At the end it kills EventStore.AdvisoryLocks and subscribes
# again. The event store uses subscription_retry_interval: 1_000 to keep the run short; with the
# default of 60 s the refusal has to last longer than about 60 s.
#
#   elixir repro.exs      # refuse for 45 s, past the lock pool's restart 30 s after it died
#   elixir repro.exs 20   # refuse for 20 s, which ends before that restart
#
# Connection settings come from PGHOST, PGPORT, PGUSER and PGPASSWORD (defaults localhost, 5432,
# postgres, postgres). REPRO_DATABASE names the database (default eventstore_blocked_lock_repro);
# the script drops and recreates it on every run.

Mix.install([
  {:eventstore, "1.4.10"},
  {:postgrex, "0.22.4"},
  {:db_connection, "2.10.2"}
])

# Postgrex logs every refused connection attempt. The script prints what it observes itself.
Logger.configure(level: :none)
# Process.info/2 cuts stacks at 8 frames by default, short of EventStore.AdvisoryLocks.handle_call/3.
:erlang.system_flag(:backtrace_depth, 20)

defmodule Store do
  use EventStore, otp_app: :repro
end

defmodule Clock do
  @moduledoc "Prints lines with the time since the outage started. Lines printed before it have no time."

  def start, do: :persistent_term.put(__MODULE__, System.monotonic_time(:millisecond))

  def elapsed_ms, do: System.monotonic_time(:millisecond) - :persistent_term.get(__MODULE__)

  def sleep_until(at_ms), do: Process.sleep(max(at_ms - elapsed_ms(), 0))

  def log(text) do
    case :persistent_term.get(__MODULE__, nil) do
      nil -> IO.puts(text)
      _started_at -> IO.puts("[+#{format_seconds(elapsed_ms())}s] #{text}")
    end
  end

  def format_seconds(milliseconds), do: :erlang.float_to_binary(milliseconds / 1000, decimals: 1)
end

defmodule Proxy do
  @moduledoc """
  TCP proxy in front of PostgreSQL. In `:forward` mode it accepts connections and relays them. In
  `:refuse` mode it closes its listening socket, so new connections fail with `:econnrefused`.
  Connections it already relays keep working in both modes.
  """

  use GenServer

  @listen_options [:binary, ip: {127, 0, 0, 1}, active: false, reuseaddr: true, packet: :raw, backlog: 128]

  def start_link(upstream_host, upstream_port) do
    upstream = {String.to_charlist(upstream_host), upstream_port}
    GenServer.start_link(__MODULE__, upstream)
  end

  @doc "The local port the proxy listens on."
  def port(proxy), do: GenServer.call(proxy, :port)

  def set_mode(proxy, mode) when mode in [:forward, :refuse], do: GenServer.call(proxy, {:set_mode, mode})

  @impl GenServer
  def init(upstream), do: {:ok, open_listener(%{upstream: upstream, port: 0, mode: :forward, listener: nil})}

  @impl GenServer
  def handle_call(:port, _from, state), do: {:reply, state.port, state}

  def handle_call({:set_mode, :refuse}, _from, state), do: {:reply, :ok, %{close_listener(state) | mode: :refuse}}

  def handle_call({:set_mode, :forward}, _from, state), do: {:reply, :ok, %{open_listener(state) | mode: :forward}}

  @impl GenServer
  def handle_info({:accepted, client}, %{mode: :forward} = state) do
    {:ok, relay} = Task.start(fn -> relay_connection(client, state.upstream) end)
    :ok = :gen_tcp.controlling_process(client, relay)
    send(relay, :go)
    {:noreply, state}
  end

  def handle_info({:accepted, client}, state) do
    :gen_tcp.close(client)
    {:noreply, state}
  end

  defp open_listener(%{listener: nil} = state) do
    {:ok, listener} = :gen_tcp.listen(state.port, @listen_options)
    {:ok, port} = :inet.port(listener)
    proxy = self()
    spawn_link(fn -> accept_connections(listener, proxy) end)
    %{state | listener: listener, port: port}
  end

  defp open_listener(state), do: state

  defp close_listener(%{listener: nil} = state), do: state

  defp close_listener(state) do
    :gen_tcp.close(state.listener)
    %{state | listener: nil}
  end

  defp accept_connections(listener, proxy) do
    case :gen_tcp.accept(listener) do
      {:ok, client} ->
        :ok = :gen_tcp.controlling_process(client, proxy)
        send(proxy, {:accepted, client})
        accept_connections(listener, proxy)

      {:error, :closed} ->
        :ok
    end
  end

  defp relay_connection(client, {host, port}) do
    receive do
      :go -> :ok
    end

    case :gen_tcp.connect(host, port, [:binary, active: false, packet: :raw]) do
      {:ok, server} ->
        start_pipe(client, server)
        start_pipe(server, client)

      {:error, _reason} ->
        :gen_tcp.close(client)
    end
  end

  defp start_pipe(from, to) do
    pipe = spawn(fn -> await_ownership_then_pipe(from, to) end)
    :ok = :gen_tcp.controlling_process(from, pipe)
    send(pipe, :go)
  end

  defp await_ownership_then_pipe(from, to) do
    receive do
      :go -> pipe(from, to)
    end
  end

  defp pipe(from, to) do
    with {:ok, data} <- :gen_tcp.recv(from, 0),
         :ok <- :gen_tcp.send(to, data) do
      pipe(from, to)
    else
      _closed ->
        :gen_tcp.close(from)
        :gen_tcp.close(to)
    end
  end
end

defmodule Database do
  @moduledoc "Database setup and the pg_locks and pg_terminate_backend queries, on a direct connection that bypasses the proxy."

  def settings do
    [
      hostname: System.get_env("PGHOST", "localhost"),
      port: String.to_integer(System.get_env("PGPORT", "5432")),
      username: System.get_env("PGUSER", "postgres"),
      password: System.get_env("PGPASSWORD", "postgres")
    ]
  end

  def connect(database) do
    {:ok, conn} = Postgrex.start_link([database: database] ++ settings())
    conn
  end

  def recreate(admin, database) do
    Postgrex.query!(admin, ~s[DROP DATABASE IF EXISTS "#{database}" WITH (FORCE)], [])
    Postgrex.query!(admin, ~s[CREATE DATABASE "#{database}"], [])
    config = Store.config([database: database] ++ settings())
    :ok = EventStore.Tasks.Init.exec(config, quiet: true)
  end

  def server_version(admin) do
    %Postgrex.Result{rows: [[version]]} =
      Postgrex.query!(admin, "SELECT split_part(current_setting('server_version'), ' ', 1)", [])

    version
  end

  @doc "Advisory locks granted in the database, as `{backend_pid, key}` ordered by key. The key is the subscription id."
  def advisory_locks(admin, database) do
    sql = """
    SELECT l.pid, l.objid
    FROM pg_locks l JOIN pg_database d ON d.oid = l.database
    WHERE l.locktype = 'advisory' AND l.granted AND d.datname = $1
    ORDER BY l.objid
    """

    %Postgrex.Result{rows: rows} = Postgrex.query!(admin, sql, [database])
    Enum.map(rows, &List.to_tuple/1)
  end

  def terminate_backend(admin, database, backend_pid) do
    sql = "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE pid = $1 AND datname = $2"
    %Postgrex.Result{rows: [[terminated?]]} = Postgrex.query!(admin, sql, [backend_pid, database])
    terminated?
  end
end

defmodule Subscriber do
  @moduledoc "Subscriber process that acknowledges every event and tells the script when it is subscribed and which events it received."

  def start(subscription_name, script), do: spawn(fn -> receive_events(subscription_name, script, nil) end)

  defp receive_events(subscription_name, script, subscription) do
    receive do
      {:subscribed, subscription} ->
        send(script, {:subscribed, subscription_name})
        receive_events(subscription_name, script, subscription)

      {:events, events} ->
        :ok = Store.ack(subscription, events)
        event_numbers = Enum.map(events, & &1.event_number)
        send(script, {:received, subscription_name, event_numbers})
        receive_events(subscription_name, script, subscription)
    end
  end
end

defmodule Watcher do
  @moduledoc """
  Prints, when it happens, every exit of a lock connection pool and of a watched subscription
  process. It learns of pool exits through `EventStore.MonitoredServer.monitor/1`, the call that
  `EventStore.AdvisoryLocks` uses for the same purpose.
  """

  alias EventStore.MonitoredServer

  def start(lock_connection) do
    script = self()
    watcher = spawn(fn -> watch_lock_pool(lock_connection, script) end)

    receive do
      {:watching, ^watcher} -> watcher
    end
  end

  def watch_subscription(watcher, subscription, subscription_name), do: send(watcher, {:watch, subscription, subscription_name})

  defp watch_lock_pool(lock_connection, script) do
    {:ok, pool_ref} = MonitoredServer.monitor(lock_connection)
    send(script, {:watching, self()})
    report_exits(pool_ref, %{})
  end

  defp report_exits(pool_ref, subscription_names) do
    receive do
      {:watch, subscription, subscription_name} ->
        ref = Process.monitor(subscription)
        report_exits(pool_ref, Map.put(subscription_names, ref, subscription_name))

      {:DOWN, ^pool_ref, :process, pool, reason} ->
        Clock.log("Lock pool #{inspect(pool)} exited: #{inspect(reason)}")
        report_exits(pool_ref, subscription_names)

      {:DOWN, ref, :process, subscription, reason} ->
        {subscription_name, subscription_names} = Map.pop!(subscription_names, ref)
        Clock.log("Subscription #{subscription_name} #{inspect(subscription)} exited: #{inspect(reason)}")
        report_exits(pool_ref, subscription_names)
    end
  end
end

defmodule Repro do
  @versioned_apps [:eventstore, :postgrex, :db_connection]
  @subscription_names ["sub-a", "sub-b"]
  @status_interval_ms 10_000
  @observation_after_refusal_ms 60_000
  @receive_window_ms 5_000
  @subscribe_timeout_ms 10_000

  # Registered names inside the event store `Store`: the lock server, and the
  # EventStore.MonitoredServer that starts and restarts its connection pool.
  @advisory_locks Module.concat([Store, EventStore.AdvisoryLocks])
  @lock_connection Module.concat([Store, EventStore.AdvisoryLocks, Postgrex])

  def run(refusal_ms) do
    database = System.get_env("REPRO_DATABASE", "eventstore_blocked_lock_repro")
    admin = Database.connect("postgres")
    Database.recreate(admin, database)
    print_versions(admin)

    settings = Database.settings()
    upstream_host = Keyword.fetch!(settings, :hostname)
    upstream_port = Keyword.fetch!(settings, :port)
    {:ok, proxy} = Proxy.start_link(upstream_host, upstream_port)
    proxy_port = Proxy.port(proxy)
    {:ok, _store} = Store.start_link(store_options(settings, database, proxy_port))
    watcher = Watcher.start(@lock_connection)

    subscriptions = Enum.map(@subscription_names, &subscribe(&1, watcher))
    Clock.log("Subscribed #{Enum.join(@subscription_names, " and ")}. Advisory locks: #{describe_locks(admin, database)}")
    append_and_report(1)

    evidence_at_ms = refusal_ms + @observation_after_refusal_ms
    status_times_ms = Enum.to_list(@status_interval_ms..(evidence_at_ms - 1)//@status_interval_ms)
    {while_refusing_ms, after_refusal_ms} = Enum.split_with(status_times_ms, &(&1 < refusal_ms))

    start_outage(proxy, admin, database, refusal_ms)
    Enum.each(while_refusing_ms, &print_status_at(&1, subscriptions, admin, database))
    Clock.sleep_until(refusal_ms)
    :ok = Proxy.set_mode(proxy, :forward)
    Clock.log("Proxy forwards new connections again")
    Enum.each(after_refusal_ms, &print_status_at(&1, subscriptions, admin, database))

    Clock.sleep_until(evidence_at_ms)
    print_evidence(subscriptions, admin, database)
    kill_advisory_locks_and_resubscribe(watcher, admin, database)
  end

  defp store_options(settings, database, proxy_port) do
    proxied = [
      hostname: "127.0.0.1",
      port: proxy_port,
      database: database,
      serializer: EventStore.TermSerializer,
      subscription_retry_interval: 1_000
    ]

    Keyword.merge(settings, proxied)
  end

  defp subscribe(subscription_name, watcher) do
    subscriber = Subscriber.start(subscription_name, self())
    {:ok, subscription} = Store.subscribe_to_all_streams(subscription_name, subscriber)
    Watcher.watch_subscription(watcher, subscription, subscription_name)
    :subscribed = await_subscribed(subscription_name)
    {subscription_name, subscription}
  end

  # A subscriber gets `{:subscribed, subscription}` only once its subscription holds the advisory lock.
  defp await_subscribed(subscription_name) do
    receive do
      {:subscribed, ^subscription_name} -> :subscribed
    after
      @subscribe_timeout_ms -> :not_subscribed
    end
  end

  defp start_outage(proxy, admin, database, refusal_ms) do
    [{lock_backend, _key_a}, {lock_backend, _key_b}] = Database.advisory_locks(admin, database)
    :ok = Proxy.set_mode(proxy, :refuse)
    Clock.start()
    terminated? = Database.terminate_backend(admin, database, lock_backend)

    Clock.log(
      "Proxy refuses new connections for #{div(refusal_ms, 1000)} s. pg_terminate_backend(#{lock_backend}), the backend holding the locks: #{terminated?}"
    )
  end

  defp print_status_at(at_ms, subscriptions, admin, database) do
    Clock.sleep_until(at_ms)
    alive_count = Enum.count(subscriptions, &subscription_alive?/1)

    Clock.log(
      "#{describe_lock_pool()} | #{describe_lock_server()} | subscriptions alive: #{alive_count} | advisory locks: #{describe_locks(admin, database)}"
    )
  end

  defp print_evidence(subscriptions, admin, database) do
    Clock.log("60 s after the proxy forwards again:")
    print_lock_server_stack()
    Clock.log("  #{describe_lock_pool()}")
    Clock.log("  The lock connection answers another process: #{query_lock_connection()}")
    Clock.log("  Advisory locks in the database: #{describe_locks(admin, database)}")
    Clock.log("  Subscription processes alive: #{Enum.map_join(subscriptions, ", ", &describe_liveness/1)}")
    append_and_report(2)
    try_new_subscription("sub-c")
    Clock.log("  #{describe_lock_server()} | advisory locks: #{describe_locks(admin, database)}")
  end

  defp print_lock_server_stack do
    advisory_locks = Process.whereis(@advisory_locks)
    {:current_stacktrace, stack} = Process.info(advisory_locks, :current_stacktrace)
    Clock.log("  #{describe_lock_server()}, stack:")
    Enum.each(stack, &IO.puts("      #{format_frame(&1)}"))
  end

  defp format_frame({module, function, arity, _location}), do: Exception.format_mfa(module, function, arity)

  # Postgrex sends this query through the same registered name that EventStore.AdvisoryLocks uses.
  defp query_lock_connection do
    {:ok, %Postgrex.Result{rows: [[backend]]}} = Postgrex.query(@lock_connection, "SELECT pg_backend_pid()", [])
    "backend #{backend}"
  end

  defp try_new_subscription(subscription_name) do
    subscriber = Subscriber.start(subscription_name, self())
    started_at = System.monotonic_time(:millisecond)
    result = subscribe_catching_exit(subscription_name, subscriber)
    elapsed_ms = System.monotonic_time(:millisecond) - started_at

    Clock.log(
      ~s[  subscribe_to_all_streams("#{subscription_name}") returned after #{Clock.format_seconds(elapsed_ms)} s: #{inspect(result, limit: 10)}]
    )
  end

  defp subscribe_catching_exit(subscription_name, subscriber) do
    Store.subscribe_to_all_streams(subscription_name, subscriber)
  catch
    :exit, reason -> {:exit, reason}
  end

  defp kill_advisory_locks_and_resubscribe(watcher, admin, database) do
    advisory_locks = Process.whereis(@advisory_locks)
    Clock.log("Killing #{inspect(@advisory_locks)} #{inspect(advisory_locks)}")
    Process.exit(advisory_locks, :kill)
    restarted = await_restart(advisory_locks)
    Clock.log("#{inspect(@advisory_locks)} restarted as #{inspect(restarted)}")

    Enum.each(@subscription_names, &resubscribe(&1, watcher))
    Clock.log("Advisory locks: #{describe_locks(admin, database)}")
    append_and_report(3)
  end

  # EventStore.Supervisor restarts every child of the instance when AdvisoryLocks dies (one_for_all).
  defp await_restart(killed) do
    case Process.whereis(@advisory_locks) do
      restarted when is_pid(restarted) and restarted != killed -> restarted
      _killed_or_none -> retry_await_restart(killed)
    end
  end

  defp retry_await_restart(killed) do
    Process.sleep(100)
    await_restart(killed)
  end

  defp resubscribe(subscription_name, watcher) do
    {^subscription_name, subscription} = subscribe(subscription_name, watcher)
    Clock.log("Subscribed #{subscription_name} again: #{inspect(subscription)}")
  end

  defp append_and_report(number) do
    event = %EventStore.EventData{event_type: "appended", data: %{number: number}, metadata: %{}}
    :ok = Store.append_to_stream("stream-1", :any_version, [event])
    received = collect_received(%{}, deadline_in(@receive_window_ms))
    Clock.log("Appended event #{number}. Received within #{div(@receive_window_ms, 1000)} s: #{describe_received(received)}")
  end

  defp collect_received(received, deadline) do
    receive do
      {:received, subscription_name, event_numbers} ->
        received
        |> Map.update(subscription_name, event_numbers, &(&1 ++ event_numbers))
        |> collect_received(deadline)
    after
      max(deadline - System.monotonic_time(:millisecond), 0) -> received
    end
  end

  defp describe_received(received) when map_size(received) == 0, do: "nothing"

  defp describe_received(received) do
    received
    |> Enum.sort()
    |> Enum.map_join(", ", &describe_events_received/1)
  end

  defp describe_events_received({subscription_name, event_numbers}), do: "#{subscription_name} #{inspect(event_numbers)}"

  defp deadline_in(timeout_ms), do: System.monotonic_time(:millisecond) + timeout_ms

  # EventStore.MonitoredServer keeps the current pool pid (nil while there is none) and a queue of
  # the messages that arrived while there was no pool, to forward to the next one it starts.
  defp describe_lock_pool do
    %{pid: pool, queue: queue} = :sys.get_state(@lock_connection)
    "lock pool #{describe_pool(pool)}, MonitoredServer queue #{:queue.len(queue)}"
  end

  defp describe_pool(nil), do: "none"
  defp describe_pool(pool), do: "#{inspect(pool)} alive=#{Process.alive?(pool)}"

  defp describe_lock_server do
    advisory_locks = Process.whereis(@advisory_locks)

    [current_function: {module, function, arity}, message_queue_len: mailbox] =
      Process.info(advisory_locks, [:current_function, :message_queue_len])

    "AdvisoryLocks in #{Exception.format_mfa(module, function, arity)}, mailbox #{mailbox}"
  end

  defp subscription_alive?({_subscription_name, subscription}), do: Process.alive?(subscription)

  defp describe_liveness({subscription_name, subscription}), do: "#{subscription_name} #{Process.alive?(subscription)}"

  defp describe_locks(admin, database) do
    case Database.advisory_locks(admin, database) do
      [] -> "none"
      locks -> describe_lock_holders(locks)
    end
  end

  defp describe_lock_holders(locks) do
    locks
    |> Enum.group_by(&elem(&1, 0), &elem(&1, 1))
    |> Enum.map_join(", ", &describe_lock_holder/1)
  end

  defp describe_lock_holder({backend, keys}), do: "backend #{backend} holds keys #{inspect(keys)}"

  defp print_versions(admin) do
    versions = Enum.map_join(@versioned_apps, ", ", &"#{&1} #{Application.spec(&1, :vsn)}")

    Clock.log(
      "#{versions}, PostgreSQL #{Database.server_version(admin)}, Elixir #{System.version()}, OTP #{System.otp_release()}"
    )
  end
end

refusal_ms =
  case System.argv() do
    [] -> 45_000
    [seconds] -> String.to_integer(seconds) * 1000
  end

Repro.run(refusal_ms)