feat(storage): record which lock owner may write a subscription's progress

lib/event_store/advisory_locks.ex

@@ -59,6 +59,28 @@ defmodule EventStore.AdvisoryLocks do
     def held_by?(%Lock{session: session}, %Session{} = current), do: session == current
   end
 
+  defmodule Lease do
+    @moduledoc false
+
+    alias EventStore.Storage.Lock.Lease, as: StorageLease
+    alias EventStore.Storage.Session
+    alias EventStore.Subscriptions.FencingToken
+
+    @type t :: %Lease{
+            ref: reference(),
+            token: FencingToken.t(),
+            last_seen: non_neg_integer() | nil,
+            session: Session.t()
+          }
+    defstruct [:ref, :token, :last_seen, :session]
+
+    def new(ref, %StorageLease{} = lease) do
+      %StorageLease{token: token, last_seen: last_seen, session: session} = lease
+
+      %Lease{ref: ref, token: token, last_seen: last_seen, session: session}
+    end
+  end
+
   defmodule Connection do
     @moduledoc false
 
@@ -79,7 +101,7 @@ defmodule EventStore.AdvisoryLocks do
   end
 
   alias DBConnection.ConnectionError
-  alias EventStore.AdvisoryLocks.{Connection, Lock, State}
+  alias EventStore.AdvisoryLocks.{Connection, Lease, Lock, State}
   alias EventStore.Storage
   alias EventStore.Storage.Session
 
@@ -131,6 +153,52 @@ defmodule EventStore.AdvisoryLocks do
     GenServer.call(server, {:try_advisory_lock, key, self(), timeout}, timeout)
   end
 
+  @doc """
+  Attempt to obtain the advisory lock of a subscription, recording the new lock
+  owner on the subscription row.
+
+  Returns `{:ok, lease}` when the lock is acquired, `{:error, :lock_already_taken}`
+  if it cannot be acquired immediately, or `{:error, :subscription_not_found}`
+  when the subscription no longer exists. The lock is released and lost the same
+  way as the one returned by `try_advisory_lock/3`, identified by `lease.ref`.
+  """
+  @spec try_subscription_lock(server :: GenServer.server(), pos_integer(), timeout()) ::
+          {:ok, Lease.t()}
+          | {:error, :lock_already_taken}
+          | {:error, :subscription_not_found}
+          | {:error, term}
+  def try_subscription_lock(server, subscription_id, timeout \\ 5_000)
+      when is_integer(subscription_id) do
+    GenServer.call(server, {:try_subscription_lock, subscription_id, self(), timeout}, timeout)
+  end
+
+  def handle_call(
+        {:try_subscription_lock, subscription_id, owner, timeout},
+        _from,
+        %State{} = state
+      ) do
+    %State{schema: schema} = state
+
+    claim =
+      query(state, fn conn ->
+        Storage.Lock.claim_subscription_lock(conn, subscription_id,
+          schema: schema,
+          timeout: timeout
+        )
+      end)
+
+    case claim do
+      {:ok, %Storage.Lock.Lease{session: session} = lease} ->
+        lock = Lock.new(subscription_id, owner, session)
+        state = monitor_acquired_lock(lock, state)
+
+        {:reply, {:ok, Lease.new(Lock.ref(lock), lease)}, state}
+
+      reply ->
+        {:reply, reply, state}
+    end
+  end
+
   def handle_call({:try_advisory_lock, key, owner, timeout}, _from, %State{} = state) do
     case try_acquire_exclusive_lock(key, owner, timeout, state) do
       {:ok, %Lock{} = lock} ->

lib/event_store/sql/init.ex

@@ -233,7 +233,8 @@ defmodule EventStore.Sql.Init do
         stream_uuid text NOT NULL,
         subscription_name text NOT NULL,
         last_seen bigint NULL,
-        created_at timestamp with time zone DEFAULT NOW() NOT NULL
+        created_at timestamp with time zone DEFAULT NOW() NOT NULL,
+        lock_generation bigint DEFAULT 0 NOT NULL
     );
     """
   end
@@ -276,7 +277,7 @@ defmodule EventStore.Sql.Init do
   defp record_event_store_schema_version do
     """
     INSERT INTO schema_migrations (major_version, minor_version, patch_version)
-    VALUES (1, 3, 2);
+    VALUES (1, 5, 0);
     """
   end
 end

lib/event_store/sql/statements.ex

@@ -21,6 +21,7 @@ defmodule EventStore.Sql.Statements do
         {:try_advisory_lock, [:schema]},
         {:try_advisory_session_lock, [:schema]},
         {:advisory_unlock, [:schema]},
+        {:claim_subscription_lock, [:schema]},
         {:subscription_ack, [:schema]},
         {:insert_snapshot, [:schema]},
         {:delete_snapshot, [:schema]},

lib/event_store/sql/statements/claim_subscription_lock.sql.eex

@@ -0,0 +1,16 @@
+WITH lock AS (
+  SELECT pg_try_advisory_lock(
+    '"<%= schema %>".subscriptions'::regclass::oid::int,
+    (CASE WHEN $1 > 2147483647 THEN mod($1, 2147483647) ELSE $1 END)::int
+  ) AS acquired, pg_backend_pid() AS pid, COALESCE(
+    NULLIF(current_setting('eventstore.session_id', true), ''),
+    set_config('eventstore.session_id', md5(random()::text || clock_timestamp()::text), false)
+  ) AS session_id
+), lease AS (
+  UPDATE "<%= schema %>".subscriptions
+  SET lock_generation = lock_generation + 1
+  WHERE subscription_id = $1 AND (SELECT acquired FROM lock)
+  RETURNING lock_generation, last_seen
+)
+SELECT lock.acquired, lock.pid, lock.session_id, lease.lock_generation, lease.last_seen
+FROM lock LEFT JOIN lease ON true;

lib/event_store/storage/lock.ex

@@ -2,7 +2,9 @@ defmodule EventStore.Storage.Lock do
   @moduledoc false
 
   alias EventStore.Sql.Statements
+  alias EventStore.Storage.Lock.Lease
   alias EventStore.Storage.Session
+  alias EventStore.Subscriptions.FencingToken
 
   @doc """
   Use two 32-bit key values for advisory locks where the first key acts as the
@@ -49,6 +51,40 @@ defmodule EventStore.Storage.Lock do
     end
   end
 
+  @doc """
+  Takes the advisory lock of a subscription and bumps its `lock_generation` in
+  the same statement, returning the lease that fences the subscription's
+  progress writes to this lock owner.
+  """
+  def claim_subscription_lock(conn, subscription_id, opts) do
+    {schema, opts} = Keyword.pop(opts, :schema)
+
+    query = Statements.claim_subscription_lock(schema)
+
+    case Postgrex.query(conn, query, [subscription_id], opts) do
+      {:ok, %Postgrex.Result{rows: [[false, _backend_pid, _session_id, _generation, _last_seen]]}} ->
+        {:error, :lock_already_taken}
+
+      {:ok, %Postgrex.Result{rows: [[true, _backend_pid, _session_id, nil, nil]]}} ->
+        with :ok <- unlock(conn, subscription_id, [schema: schema] ++ opts) do
+          {:error, :subscription_not_found}
+        end
+
+      {:ok, %Postgrex.Result{rows: [[true, backend_pid, session_id, generation, last_seen]]}} ->
+        token = FencingToken.new(subscription_id, generation)
+
+        {:ok, Lease.new(token, last_seen, Session.new(backend_pid, session_id))}
+
+      {:error, %Postgrex.Error{} = error} ->
+        _ = unlock(conn, subscription_id, [schema: schema] ++ opts)
+
+        {:error, error}
+
+      {:error, %DBConnection.ConnectionError{}} = reply ->
+        reply
+    end
+  end
+
   def unlock(conn, key, opts) do
     {schema, opts} = Keyword.pop(opts, :schema)
 

lib/event_store/storage/lock/lease.ex

@@ -0,0 +1,19 @@
+defmodule EventStore.Storage.Lock.Lease do
+  @moduledoc false
+
+  alias EventStore.Storage.Session
+  alias EventStore.Subscriptions.FencingToken
+
+  @type t :: %__MODULE__{
+          token: FencingToken.t(),
+          last_seen: non_neg_integer() | nil,
+          session: Session.t()
+        }
+
+  @enforce_keys [:token, :last_seen, :session]
+  defstruct [:token, :last_seen, :session]
+
+  def new(%FencingToken{} = token, last_seen, %Session{} = session) do
+    %__MODULE__{token: token, last_seen: last_seen, session: session}
+  end
+end

lib/event_store/subscriptions/fencing_token.ex

@@ -0,0 +1,16 @@
+defmodule EventStore.Subscriptions.FencingToken do
+  @moduledoc false
+
+  @type t :: %__MODULE__{
+          subscription_id: pos_integer(),
+          generation: pos_integer()
+        }
+
+  @enforce_keys [:subscription_id, :generation]
+  defstruct [:subscription_id, :generation]
+
+  def new(subscription_id, generation)
+      when is_integer(subscription_id) and is_integer(generation) and generation > 0 do
+    %__MODULE__{subscription_id: subscription_id, generation: generation}
+  end
+end

lib/event_store/tasks/migrations.ex

@@ -19,7 +19,8 @@ defmodule EventStore.Tasks.Migrations do
     "1.1.0",
     "1.2.0",
     "1.3.0",
-    "1.3.2"
+    "1.3.2",
+    "1.5.0"
   ]
 
   @dialyzer {:no_return, exec: 2, handle_response: 1}

priv/event_store/migrations/v1.5.0.sql

@@ -0,0 +1,8 @@
+DO $migration$
+BEGIN
+  ALTER TABLE subscriptions ADD COLUMN lock_generation bigint NOT NULL DEFAULT 0;
+
+  -- Record schema migration
+  INSERT INTO schema_migrations (major_version, minor_version, patch_version) VALUES (1, 5, 0);
+END;
+$migration$ LANGUAGE plpgsql;

test/advisory_locks_test.exs

@@ -42,6 +42,84 @@ defmodule EventStore.AdvisoryLocksTest do
     end
   end
 
+  describe "acquire subscription lock" do
+    test "should acquire the lock and record the lock owner", context do
+      %{subscription_id: subscription_id} = create_subscription(context)
+
+      assert {:ok, %AdvisoryLocks.Lease{} = lease} =
+               AdvisoryLocks.try_subscription_lock(@locks, subscription_id)
+
+      assert is_reference(lease.ref)
+      assert lease.token == EventStore.Subscriptions.FencingToken.new(subscription_id, 1)
+      assert lease.last_seen == nil
+    end
+
+    test "should record a new lock owner each time the lock is acquired", context do
+      %{subscription_id: subscription_id} = create_subscription(context)
+
+      assert {:ok, %{token: %{generation: 1}}} =
+               AdvisoryLocks.try_subscription_lock(@locks, subscription_id)
+
+      assert {:ok, %{token: %{generation: 2}}} =
+               AdvisoryLocks.try_subscription_lock(@locks, subscription_id)
+    end
+
+    test "should fail to acquire the lock when another session holds it", context do
+      %{conn: conn, schema: schema} = context
+      %{subscription_id: subscription_id} = create_subscription(context)
+
+      assert {:ok, _lease} =
+               Storage.Lock.claim_subscription_lock(conn, subscription_id, schema: schema)
+
+      assert {:error, :lock_already_taken} =
+               AdvisoryLocks.try_subscription_lock(@locks, subscription_id)
+    end
+
+    test "should fail to acquire the lock of a subscription that does not exist" do
+      assert {:error, :subscription_not_found} =
+               AdvisoryLocks.try_subscription_lock(@locks, 999_999)
+    end
+
+    test "should send `lock_released` message with the lease reference", context do
+      %{subscription_id: subscription_id} = create_subscription(context)
+
+      {:ok, %{ref: ref}} = AdvisoryLocks.try_subscription_lock(@locks, subscription_id)
+
+      shutdown_database_connection()
+
+      assert_receive({AdvisoryLocks, :lock_released, ^ref, :shutdown})
+    end
+
+    test "should release the lock when the owner terminates", context do
+      %{conn: conn, schema: schema} = context
+      %{subscription_id: subscription_id} = create_subscription(context)
+      reply_to = self()
+
+      pid =
+        spawn_link(fn ->
+          {:ok, _lease} = AdvisoryLocks.try_subscription_lock(@locks, subscription_id)
+
+          send(reply_to, :lock_acquired)
+
+          receive do
+            :shutdown -> :ok
+          end
+        end)
+
+      assert_receive :lock_acquired
+
+      assert {:error, :lock_already_taken} =
+               Storage.Lock.try_acquire_exclusive_lock(conn, subscription_id, schema: schema)
+
+      send(pid, :shutdown)
+
+      Wait.until(fn ->
+        assert :ok =
+                 Storage.Lock.try_acquire_exclusive_lock(conn, subscription_id, schema: schema)
+      end)
+    end
+  end
+
   describe "release lock" do
     test "should release lock when process terminates", %{conn: conn, schema: schema} do
       reply_to = self()
@@ -513,4 +591,11 @@ defmodule EventStore.AdvisoryLocksTest do
   end
 
   defp lock_pool_pid, do: Process.whereis(@conn)
+
+  defp create_subscription(%{conn: conn, schema: schema}) do
+    {:ok, %Storage.Subscription{} = subscription} =
+      Storage.subscribe_to_stream(conn, "$all", "test_subscription", schema: schema)
+
+    subscription
+  end
 end

test/list_event_store_migrations_test.exs

@@ -26,6 +26,7 @@ defmodule EventStore.ListEventStoreMigrationsTest do
         1.2.0     \tpending
         1.3.0     \tpending
         1.3.2     \tpending
+        1.5.0     \tpending
 
       """
 

test/migrate_event_store_test.exs

@@ -88,6 +88,7 @@ defmodule EventStore.MigrateEventStoreTest do
       "1\t1\t0\tNOW" -> true
       "1\t2\t0\tNOW" -> true
       "1\t3\t0\tNOW" -> true
+      "1\t3\t2\tNOW" -> true
       _line -> false
     end)
   end