fix(subscriptions): release advisory locks when the lock connection reconnects on a new session

lib/event_store/advisory_locks.ex

@@ -7,11 +7,24 @@ defmodule EventStore.AdvisoryLocks do
 
   use GenServer
 
+  require Logger
+
   alias EventStore.MonitoredServer
 
+  @verify_retry 1_000
+
   defmodule State do
     @moduledoc false
-    defstruct [:conn, :query_timeout, :ref, :schema, locks: %{}]
+    defstruct [
+      :conn,
+      :connection,
+      :query_timeout,
+      :ref,
+      :schema,
+      :verify_timer,
+      locks: %{},
+      released: []
+    ]
 
     def new(opts) do
       conn = Keyword.fetch!(opts, :conn)
@@ -25,22 +38,47 @@ defmodule EventStore.AdvisoryLocks do
   defmodule Lock do
     @moduledoc false
 
+    alias EventStore.Storage.Session
+
     @type t :: %Lock{
             key: non_neg_integer(),
             owner: pid(),
-            ref: reference()
+            ref: reference(),
+            session: Session.t()
           }
-    defstruct [:key, :owner, :ref]
+    defstruct [:key, :owner, :ref, :session]
 
-    def new(key, owner) do
-      %Lock{key: key, owner: owner, ref: make_ref()}
+    def new(key, owner, %Session{} = session) do
+      %Lock{key: key, owner: owner, ref: make_ref(), session: session}
     end
 
     def ref(%Lock{ref: ref}), do: ref
+
+    def held_by?(%Lock{session: session}, %Session{} = current), do: session == current
+  end
+
+  defmodule Connection do
+    @moduledoc false
+
+    @type t :: %Connection{pid: pid(), monitor_ref: reference()}
+    defstruct [:pid, :monitor_ref]
+
+    def monitor(pid), do: %Connection{pid: pid, monitor_ref: Process.monitor(pid)}
+
+    def demonitor(nil), do: :ok
+
+    def demonitor(%Connection{monitor_ref: monitor_ref}),
+      do: Process.demonitor(monitor_ref, [:flush])
+
+    # Without a known connection any notice may be about the one holding the locks.
+    def current?(nil, _pid), do: true
+    def current?(%Connection{pid: pid}, pid), do: true
+    def current?(%Connection{}, _pid), do: false
   end
 
-  alias EventStore.AdvisoryLocks.{Lock, State}
+  alias EventStore.AdvisoryLocks.{Connection, Lock, State}
   alias EventStore.Storage
+  alias EventStore.Storage.Session
 
   def start_link(opts) do
     {start_opts, advisory_locks_opts} =
@@ -58,6 +96,10 @@ defmodule EventStore.AdvisoryLocks do
 
     {:ok, ref} = MonitoredServer.monitor(conn)
 
+    # A connection made before this server was registered announced itself to
+    # nobody, so it is replaced by one whose notices can be received.
+    :ok = DBConnection.disconnect_all(conn, 0)
+
     {:ok, %State{state | ref: ref}}
   end
 
@@ -101,9 +143,12 @@ defmodule EventStore.AdvisoryLocks do
   defp try_acquire_exclusive_lock(key, owner, timeout, %State{} = state) do
     %State{conn: conn, schema: schema} = state
 
-    with :ok <-
-           Storage.Lock.try_acquire_exclusive_lock(conn, key, schema: schema, timeout: timeout) do
-      {:ok, Lock.new(key, owner)}
+    with {:ok, %Session{} = session} <-
+           Storage.Lock.try_acquire_exclusive_session_lock(conn, key,
+             schema: schema,
+             timeout: timeout
+           ) do
+      {:ok, Lock.new(key, owner, session)}
     end
   end
 
@@ -121,9 +166,34 @@ defmodule EventStore.AdvisoryLocks do
 
     notify_lost_locks(locks, reason)
 
-    {:noreply, %State{state | locks: %{}}}
+    {:noreply, %State{state | connection: nil, locks: %{}, released: [], verify_timer: nil}}
   end
 
+  # A session that could not be verified never reported a disconnect, so the
+  # pending unlocks are reconciled again once it can answer.
+  def handle_info(:verify_sessions, %State{} = state),
+    do: {:noreply, verify_sessions(%State{state | verify_timer: nil})}
+
+  # The connection pool outlives its connections, so a reconnect replaces the
+  # session that held the locks without the pool going down.
+  def handle_info({:connected, conn_pid}, %State{} = state) when is_pid(conn_pid),
+    do: {:noreply, connected(conn_pid, state)}
+
+  def handle_info({:connected, conn_pid, _tag}, %State{} = state) when is_pid(conn_pid),
+    do: {:noreply, connected(conn_pid, state)}
+
+  def handle_info({:disconnected, conn_pid}, %State{} = state) when is_pid(conn_pid),
+    do: {:noreply, disconnected(conn_pid, :session_ended, state)}
+
+  def handle_info({:disconnected, conn_pid, _tag}, %State{} = state) when is_pid(conn_pid),
+    do: {:noreply, disconnected(conn_pid, :session_ended, state)}
+
+  def handle_info(
+        {:DOWN, monitor_ref, :process, conn_pid, reason},
+        %State{connection: %Connection{monitor_ref: monitor_ref}} = state
+      ),
+      do: {:noreply, disconnected(conn_pid, reason, state)}
+
   # Release lock when the lock owner process terminates.
   def handle_info({:DOWN, ref, :process, _pid, _reason}, %State{} = state) do
     %State{locks: locks} = state
@@ -150,10 +220,80 @@ defmodule EventStore.AdvisoryLocks do
     {:stop, reason, state}
   end
 
+  defp connected(conn_pid, %State{connection: connection} = state) do
+    Connection.demonitor(connection)
+
+    verify_sessions(%State{state | connection: Connection.monitor(conn_pid)})
+  end
+
+  # 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
+    if Connection.current?(connection, conn_pid) do
+      %State{locks: locks, released: released} = state
+
+      Connection.demonitor(connection)
+
+      notify_lost_locks(locks, reason)
+
+      %State{state | connection: nil, locks: %{}, released: Map.values(locks) ++ released}
+    else
+      state
+    end
+  end
+
+  defp verify_sessions(%State{locks: locks, released: []} = state) when map_size(locks) == 0,
+    do: state
+
+  defp verify_sessions(%State{} = state) do
+    %State{conn: conn, locks: locks, query_timeout: query_timeout, released: released} = state
+
+    case Session.current(conn, timeout: query_timeout) do
+      {:ok, %Session{} = session} ->
+        {held, lost} = Map.split_with(locks, fn {_ref, lock} -> Lock.held_by?(lock, session) end)
+
+        notify_lost_locks(lost, :session_ended)
+
+        # A lock whose session has ended died with it, but one still held by this
+        # session is only cleared once its unlock is acknowledged.
+        pending =
+          for %Lock{key: key} = lock <- released,
+              Lock.held_by?(lock, session),
+              release_lock(key, state) != :ok,
+              do: lock
+
+        # A live session confirms the locks still held, so polling stops once
+        # the pending unlocks are drained.
+        maybe_schedule_verify_retry(%State{state | locks: held, released: pending})
+
+      # A failed query does not prove the session ended, so verification is
+      # retried until a live session can confirm the locks and release them.
+      {:error, _error} ->
+        schedule_verify_retry(state)
+    end
+  end
+
+  defp maybe_schedule_verify_retry(%State{released: []} = state), do: state
+  defp maybe_schedule_verify_retry(%State{} = state), do: schedule_verify_retry(state)
+
+  defp schedule_verify_retry(%State{verify_timer: ref} = state) when is_reference(ref), do: state
+
+  defp schedule_verify_retry(%State{} = state),
+    do: %State{state | verify_timer: Process.send_after(self(), :verify_sessions, @verify_retry)}
+
+  defp notify_lost_locks(locks, _reason) when map_size(locks) == 0, do: :ok
+
   defp notify_lost_locks(locks, reason) do
-    for {_ref, %Lock{} = lock} <- locks do
+    Logger.warning(
+      "Lost #{map_size(locks)} advisory lock(s) held by this event store due to: " <>
+        inspect(reason)
+    )
+
+    for {owner_ref, %Lock{} = lock} <- locks do
       %Lock{owner: owner, ref: ref} = lock
 
+      Process.demonitor(owner_ref, [:flush])
+
       :ok = Process.send(owner, {__MODULE__, :lock_released, ref, reason}, [])
     end
 

lib/event_store/config.ex

@@ -196,6 +196,23 @@ defmodule EventStore.Config do
     |> Keyword.put(:pool_size, 1)
   end
 
+  @doc """
+  Add a connection listener while keeping any listeners already configured.
+
+  A listener given as a registered name is wrapped in a `{name, node}` tuple, so
+  notifications sent before the name is registered are dropped instead of
+  crashing the connection.
+  """
+  def add_connection_listener(opts, name) when is_atom(name) do
+    listener = {name, node()}
+
+    case Keyword.get(opts, :connection_listeners) do
+      nil -> Keyword.put(opts, :connection_listeners, [listener])
+      {listeners, tag} -> Keyword.put(opts, :connection_listeners, {[listener | listeners], tag})
+      listeners -> Keyword.put(opts, :connection_listeners, [listener | listeners])
+    end
+  end
+
   # 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/sql/statements.ex

@@ -19,11 +19,13 @@ defmodule EventStore.Sql.Statements do
         {:insert_subscription, [:schema]},
         {:delete_subscription, [:schema]},
         {:try_advisory_lock, [:schema]},
+        {:try_advisory_session_lock, [:schema]},
         {:advisory_unlock, [:schema]},
         {:subscription_ack, [:schema]},
         {:insert_snapshot, [:schema]},
         {:delete_snapshot, [:schema]},
         {:query_all_subscriptions, [:schema]},
+        {:query_session, []},
         {:query_snapshot, [:schema]},
         {:query_stream_info, [:schema]},
         {:query_stream_events_backward, [:schema]},

lib/event_store/sql/statements/query_session.sql.eex

@@ -0,0 +1,4 @@
+SELECT pg_backend_pid(), COALESCE(
+  NULLIF(current_setting('eventstore.session_id', true), ''),
+  set_config('eventstore.session_id', md5(random()::text || clock_timestamp()::text), false)
+);

lib/event_store/sql/statements/try_advisory_session_lock.sql.eex

@@ -0,0 +1,7 @@
+SELECT pg_try_advisory_lock(
+  '"<%= schema %>".subscriptions'::regclass::oid::int,
+  (CASE WHEN $1 > 2147483647 THEN mod($1, 2147483647) ELSE $1 END)::int
+), pg_backend_pid(), COALESCE(
+  NULLIF(current_setting('eventstore.session_id', true), ''),
+  set_config('eventstore.session_id', md5(random()::text || clock_timestamp()::text), false)
+);

lib/event_store/storage/lock.ex

@@ -2,6 +2,7 @@ defmodule EventStore.Storage.Lock do
   @moduledoc false
 
   alias EventStore.Sql.Statements
+  alias EventStore.Storage.Session
 
   @doc """
   Use two 32-bit key values for advisory locks where the first key acts as the
@@ -27,6 +28,27 @@ defmodule EventStore.Storage.Lock do
     end
   end
 
+  @doc """
+  Like `try_acquire_exclusive_lock/3`, but also returns the session holding the
+  lock, which releases it when it ends.
+  """
+  def try_acquire_exclusive_session_lock(conn, key, opts) do
+    {schema, opts} = Keyword.pop(opts, :schema)
+
+    query = Statements.try_advisory_session_lock(schema)
+
+    case Postgrex.query(conn, query, [key], opts) do
+      {:ok, %Postgrex.Result{rows: [[true, backend_pid, session_id]]}} ->
+        {:ok, Session.new(backend_pid, session_id)}
+
+      {:ok, %Postgrex.Result{rows: [[false, _backend_pid, _session_id]]}} ->
+        {:error, :lock_already_taken}
+
+      {:error, _error} = reply ->
+        reply
+    end
+  end
+
   def unlock(conn, key, opts) do
     {schema, opts} = Keyword.pop(opts, :schema)
 

lib/event_store/storage/session.ex

@@ -0,0 +1,26 @@
+defmodule EventStore.Storage.Session do
+  @moduledoc false
+
+  # The id is a session level setting the session sets the first time it is read,
+  # because pg_stat_activity hides backend_start from a role switched to with SET ROLE.
+
+  alias EventStore.Sql.Statements
+
+  @type t :: %__MODULE__{backend_pid: pos_integer(), id: String.t()}
+
+  defstruct [:backend_pid, :id]
+
+  def new(backend_pid, id) when is_integer(backend_pid) and is_binary(id) and id != "" do
+    %__MODULE__{backend_pid: backend_pid, id: id}
+  end
+
+  def current(conn, opts) do
+    case Postgrex.query(conn, Statements.query_session(), [], opts) do
+      {:ok, %Postgrex.Result{rows: [[backend_pid, id]]}} ->
+        {:ok, new(backend_pid, id)}
+
+      {:error, _error} = reply ->
+        reply
+    end
+  end
+end

lib/event_store/supervisor.ex

@@ -58,7 +58,11 @@ defmodule EventStore.Supervisor do
 
         advisory_locks_name = Module.concat([name, AdvisoryLocks])
         advisory_locks_postgrex_conn = Module.concat([advisory_locks_name, Postgrex])
-        advisory_locks_config = Config.advisory_locks_postgrex_opts(config)
+
+        advisory_locks_config =
+          config
+          |> Config.advisory_locks_postgrex_opts()
+          |> Config.add_connection_listener(advisory_locks_name)
 
         subscriptions_name = Module.concat([name, Subscriptions.Supervisor])
         subscriptions_registry_name = Module.concat([name, Subscriptions.Registry])