fix(subscriptions): tell lock owners their lock is gone before the connection finishes closing

lib/event_store/advisory_locks.ex

@@ -18,6 +18,7 @@ defmodule EventStore.AdvisoryLocks do
     defstruct [
       :conn,
       :connection,
+      :lost_connection,
       :monitor,
       :query_timeout,
       :ref,
@@ -320,6 +321,11 @@ defmodule EventStore.AdvisoryLocks do
   def handle_info({:connected, conn_pid, _tag}, %State{} = state) when is_pid(conn_pid),
     do: {:noreply, connected(conn_pid, state)}
 
+  # Sent as the connection starts closing, which can take a long time when the
+  # database stops answering, ahead of its `:disconnected` notice.
+  def handle_info({__MODULE__, :session_lost, conn_pid}, %State{} = state) when is_pid(conn_pid),
+    do: {:noreply, disconnected(conn_pid, :session_ended, state)}
+
   def handle_info({:disconnected, conn_pid}, %State{} = state) when is_pid(conn_pid),
     do: {:noreply, disconnected(conn_pid, :session_ended, state)}
 
@@ -361,9 +367,15 @@ defmodule EventStore.AdvisoryLocks do
   defp connected(conn_pid, %State{connection: connection} = state) do
     Connection.demonitor(connection)
 
-    verify_sessions(%State{state | connection: Connection.monitor(conn_pid)})
+    verify_sessions(%State{
+      state
+      | connection: Connection.monitor(conn_pid),
+        lost_connection: nil
+    })
   end
 
+  defp disconnected(conn_pid, _reason, %State{lost_connection: conn_pid} = state), do: state
+
   # A lock taken through the next connection can be recorded before this notice
   # is handled, so it is released here too and only unlocked once verified.
   defp disconnected(conn_pid, reason, %State{connection: connection} = state) do
@@ -374,7 +386,13 @@ defmodule EventStore.AdvisoryLocks do
 
       notify_lost_locks(locks, reason)
 
-      %State{state | connection: nil, locks: %{}, released: Map.values(locks) ++ released}
+      %State{
+        state
+        | connection: nil,
+          lost_connection: conn_pid,
+          locks: %{},
+          released: Map.values(locks) ++ released
+      }
     else
       state
     end

lib/event_store/advisory_locks/protocol.ex

@@ -0,0 +1,75 @@
+defmodule EventStore.AdvisoryLocks.Protocol do
+  @moduledoc false
+
+  # Postgrex closes a connection by sending a cancel request first, and both its
+  # connect and its wait for the reply use the :timeout option, so this can take
+  # that long when the database stops answering.
+  # The lock owners cannot wait that long to hear the session is gone.
+
+  @behaviour DBConnection
+
+  @session_listener {__MODULE__, :session_listener}
+
+  def start_link(opts) do
+    opts = Postgrex.Utils.default_opts(opts)
+
+    DBConnection.start_link(__MODULE__, opts)
+  end
+
+  @impl DBConnection
+  def connect(opts) do
+    Process.put(@session_listener, Keyword.get(opts, :session_listener))
+
+    Postgrex.Protocol.connect(opts)
+  end
+
+  @impl DBConnection
+  def disconnect(error, state) do
+    notify_session_lost(error)
+
+    Postgrex.Protocol.disconnect(error, state)
+  end
+
+  defp notify_session_lost(_error) do
+    case Process.get(@session_listener) do
+      nil -> :ok
+      listener -> send(listener, {EventStore.AdvisoryLocks, :session_lost, self()})
+    end
+  end
+
+  @impl DBConnection
+  defdelegate checkout(state), to: Postgrex.Protocol
+
+  @impl DBConnection
+  defdelegate ping(state), to: Postgrex.Protocol
+
+  @impl DBConnection
+  defdelegate handle_begin(opts, state), to: Postgrex.Protocol
+
+  @impl DBConnection
+  defdelegate handle_commit(opts, state), to: Postgrex.Protocol
+
+  @impl DBConnection
+  defdelegate handle_rollback(opts, state), to: Postgrex.Protocol
+
+  @impl DBConnection
+  defdelegate handle_status(opts, state), to: Postgrex.Protocol
+
+  @impl DBConnection
+  defdelegate handle_prepare(query, opts, state), to: Postgrex.Protocol
+
+  @impl DBConnection
+  defdelegate handle_execute(query, params, opts, state), to: Postgrex.Protocol
+
+  @impl DBConnection
+  defdelegate handle_close(query, opts, state), to: Postgrex.Protocol
+
+  @impl DBConnection
+  defdelegate handle_declare(query, params, opts, state), to: Postgrex.Protocol
+
+  @impl DBConnection
+  defdelegate handle_fetch(query, cursor, opts, state), to: Postgrex.Protocol
+
+  @impl DBConnection
+  defdelegate handle_deallocate(query, cursor, opts, state), to: Postgrex.Protocol
+end

lib/event_store/config.ex

@@ -213,6 +213,15 @@ defmodule EventStore.Config do
     end
   end
 
+  @doc """
+  Tell the advisory locks process as soon as its connection starts closing.
+
+  The listener is wrapped in a `{name, node}` tuple, so a notice sent before the
+  name is registered is dropped instead of crashing the connection.
+  """
+  def add_session_listener(opts, name) when is_atom(name),
+    do: Keyword.put(opts, :session_listener, {name, node()})
+
   # Get the optional session mode pool to be used for persistent Postgres
   # connections (advisory locks and notifications). Uses the default Postgres
   # configuration if not specified.

lib/event_store/supervisor.ex

@@ -67,6 +67,7 @@ defmodule EventStore.Supervisor do
           |> Config.advisory_locks_postgrex_opts()
           |> Keyword.put(:name, advisory_locks_postgrex_conn)
           |> Config.add_connection_listener(advisory_locks_name)
+          |> Config.add_session_listener(advisory_locks_name)
 
         subscriptions_name = Module.concat([name, Subscriptions.Supervisor])
         subscriptions_registry_name = Module.concat([name, Subscriptions.Registry])
@@ -84,7 +85,7 @@ defmodule EventStore.Supervisor do
             ),
             Supervisor.child_spec(
               {MonitoredServer,
-               mfa: {Postgrex, :start_link, [advisory_locks_config]},
+               mfa: {AdvisoryLocks.Protocol, :start_link, [advisory_locks_config]},
                name: advisory_locks_postgrex_monitor,
                backoff_min: 30_000},
               id: advisory_locks_postgrex_monitor

test/advisory_locks_test.exs

@@ -155,7 +155,7 @@ defmodule EventStore.AdvisoryLocksTest do
 
       shutdown_database_connection()
 
-      assert_receive({AdvisoryLocks, :lock_released, ^ref, :shutdown})
+      assert_receive({AdvisoryLocks, :lock_released, ^ref, :session_ended})
     end
 
     test "should release the lock when the owner terminates", context do
@@ -345,7 +345,7 @@ defmodule EventStore.AdvisoryLocksTest do
 
       shutdown_database_connection()
 
-      assert_receive({AdvisoryLocks, :lock_released, ^lock, :shutdown})
+      assert_receive({AdvisoryLocks, :lock_released, ^lock, _reason})
     end
 
     test "should send `lock_released` message when the lock session ends and the connection reconnects",
@@ -406,6 +406,89 @@ defmodule EventStore.AdvisoryLocksTest do
       refute_receive {AdvisoryLocks, :lock_released, ^lock, _reason}, 500
     end
 
+    test "should send `lock_released` message when the database stops answering new connections",
+         %{conn: conn, schema: schema} do
+      name = :proxied_lock_store
+      locks = Module.concat([name, AdvisoryLocks])
+      lock_pool = Module.concat([locks, Postgrex])
+
+      config = Config.parsed(TestEventStore, :eventstore)
+
+      proxy =
+        start_supervised!(
+          {EventStore.TcpProxy,
+           {Keyword.fetch!(config, :hostname), Keyword.get(config, :port, 5432)}}
+        )
+
+      start_supervised!(
+        {TestEventStore,
+         name: name,
+         schema: schema,
+         session_mode_pool: [
+           port: EventStore.TcpProxy.port(proxy),
+           timeout: 60_000,
+           idle_interval: 100
+         ]}
+      )
+
+      try do
+        {:ok, lock} = AdvisoryLocks.try_advisory_lock(locks, 1)
+
+        %Postgrex.Result{rows: [[backend_pid]]} =
+          Postgrex.query!(lock_pool, "SELECT pg_backend_pid()", [])
+
+        :ok = EventStore.TcpProxy.blackhole(proxy)
+
+        %Postgrex.Result{rows: [[true]]} =
+          Postgrex.query!(conn, "SELECT pg_terminate_backend($1)", [backend_pid])
+
+        Wait.until(fn ->
+          assert %Postgrex.Result{rows: [[0]]} =
+                   Postgrex.query!(
+                     conn,
+                     "SELECT count(*)::int FROM pg_stat_activity WHERE pid = $1",
+                     [backend_pid]
+                   )
+        end)
+
+        assert :ok = Storage.Lock.try_acquire_exclusive_lock(conn, 1, schema: schema)
+
+        assert_receive {AdvisoryLocks, :lock_released, ^lock, _reason}, 1_000
+      after
+        stop_supervised!(EventStore.TcpProxy)
+      end
+    end
+
+    test "should keep a lock acquired while the lost connection is still closing and unlock it when its owner exits",
+         %{conn: conn, schema: schema} do
+      lock_connection = lock_connection_pid()
+
+      send(@locks, {AdvisoryLocks, :session_lost, lock_connection})
+
+      owner = start_lock_owner(1)
+
+      assert_receive {:owner_acquired, ^owner, {:ok, lock}}, 1_000
+
+      send(@locks, {:disconnected, lock_connection})
+
+      # The lock server has handled the notice once it answers, and the owner forwards
+      # its messages in order, so the probe arrives after any notification the notice sent.
+      :sys.get_state(@locks)
+      send(owner, :probe)
+      assert_receive {:owner_notified, ^owner, :probe}, 1_000
+
+      refute_received {:owner_notified, ^owner, {AdvisoryLocks, :lock_released, ^lock, _reason}}
+
+      assert {:error, :lock_already_taken} =
+               Storage.Lock.try_acquire_exclusive_lock(conn, 1, schema: schema)
+
+      Process.exit(owner, :kill)
+
+      Wait.until(fn ->
+        assert :ok = Storage.Lock.try_acquire_exclusive_lock(conn, 1, schema: schema)
+      end)
+    end
+
     test "should unlock a lock reported lost while its session is still alive",
          %{conn: conn, schema: schema} do
       {:ok, lock} = AdvisoryLocks.try_advisory_lock(@locks, 1)
@@ -812,6 +895,25 @@ defmodule EventStore.AdvisoryLocksTest do
     end
   end
 
+  defp start_lock_owner(key) do
+    test_process = self()
+
+    spawn(fn ->
+      send(test_process, {:owner_acquired, self(), AdvisoryLocks.try_advisory_lock(@locks, key)})
+
+      forward_to_test_process(test_process)
+    end)
+  end
+
+  defp forward_to_test_process(test_process) do
+    receive do
+      message ->
+        send(test_process, {:owner_notified, self(), message})
+
+        forward_to_test_process(test_process)
+    end
+  end
+
   defp hold_only_connection(conn \\ lock_pool_pid()) do
     reply_to = self()