Only resolve lazy allowances if absolutely necessary

lib/nimble_ownership.ex

@@ -405,10 +405,7 @@ defmodule NimbleOwnership do
     allowances: %{},
 
     # This is used to track which PIDs we're monitoring, to avoid double-monitoring.
-    monitored_pids: MapSet.new(),
-
-    # This boolean field tracks whether there are any lazy calls in the allowances.
-    lazy_calls: false
+    monitored_pids: MapSet.new()
   ]
 
   ## Callbacks
@@ -468,15 +465,12 @@ defmodule NimbleOwnership do
           state
           |> maybe_monitor_pid(pid_with_access)
           |> put_in([Access.key!(:allowances), Access.key(pid_to_allow, %{}), key], owner_pid)
-          |> update_in([Access.key!(:lazy_calls)], &(&1 or is_function(pid_to_allow, 0)))
 
         {:reply, :ok, state}
     end
   end
 
   def handle_call({:get_and_update, owner_pid, key, fun}, _from, %__MODULE__{} = state) do
-    state = revalidate_lazy_calls(state)
-
     case state.mode do
       {:shared, shared_owner_pid} when shared_owner_pid != owner_pid ->
         error = %Error{key: key, reason: {:not_shared_owner, shared_owner_pid}}
@@ -486,6 +480,8 @@ defmodule NimbleOwnership do
         :ok
     end
 
+    state = resolve_lazy_calls_for_key(state, key)
+
     if other_owner = state.allowances[owner_pid][key] do
       throw({:reply, {:error, %Error{key: key, reason: {:already_allowed, other_owner}}}, state})
     end
@@ -526,15 +522,15 @@ defmodule NimbleOwnership do
   end
 
   def handle_call({:fetch_owner, callers, key}, _from, %__MODULE__{mode: :private} = state) do
-    state = revalidate_lazy_calls(state)
-
-    Enum.find_value(callers, {:reply, :error, state}, fn caller ->
-      case state do
-        %{allowances: %{^caller => %{^key => owner_pid}}} -> {:reply, {:ok, owner_pid}, state}
-        %{owners: %{^caller => %{^key => _meta}}} -> {:reply, {:ok, caller}, state}
-        _ -> nil
-      end
-    end)
+    with nil <- fetch_owner_once(state, callers, key) do
+      state = resolve_lazy_calls_for_key(state, key)
+      {fetch_owner_once(state, callers, key), state}
+    end
+    |> case do
+      {nil, new_state} -> {:reply, :error, new_state}
+      {found, new_state} -> {:reply, {:ok, found}, new_state}
+      found -> {:reply, {:ok, found}, state}
+    end
   end
 
   def handle_call({:get_owned, owner_pid, default}, _from, %__MODULE__{} = state) do
@@ -616,51 +612,39 @@ defmodule NimbleOwnership do
     end
   end
 
-  defp revalidate_lazy_calls(state) do
-    state.allowances
-    |> Enum.reduce({[], [], false}, fn
-      {key, value}, {result, resolved, unresolved} when is_function(key, 0) ->
-        resolve_once(key.(), {key, value}, {result, resolved, unresolved})
-
-      kv, {result, resolved, unresolved} ->
-        {[kv | result], resolved, unresolved}
+  defp fetch_owner_once(state, callers, key) do
+    Enum.find_value(callers, fn caller ->
+      case state do
+        %{owners: %{^caller => %{^key => _meta}}} -> caller
+        %{allowances: %{^caller => %{^key => owner_pid}}} -> owner_pid
+        _ -> nil
+      end
     end)
-    |> fix_resolved(state)
   end
 
-  defp resolve_once(pid, {key, value}, {result, resolved, unresolved}) when is_pid(pid) do
-    {[{pid, value} | result], [{key, pid} | resolved], unresolved}
-  end
+  defp resolve_lazy_calls_for_key(state, key) do
+    updated_allowances =
+      state.allowances
+      |> Enum.reduce(state.allowances, fn
+        {fun, value}, allowances when is_function(fun, 0) and is_map_key(value, key) ->
+          result =
+            fun.()
+            |> List.wrap()
+            |> Enum.group_by(&is_pid/1)
 
-  defp resolve_once([pid | pids], {key, value}, {result, resolved, unresolved})
-       when is_pid(pid) do
-    resolve_once(
-      pids,
-      {key, value},
-      {[{pid, value} | result], [{key, pid} | resolved], unresolved}
-    )
-  end
+          allowances =
+            result
+            |> Map.get(true, [])
+            |> Enum.reduce(allowances, fn pid, allowances ->
+              Map.update(allowances, pid, value, &Map.merge(&1, value))
+            end)
 
-  defp resolve_once([_not_a_pid | pids], kv, {result, resolved, _unresolved}) do
-    resolve_once(pids, kv, {[kv | result], resolved, true})
-  end
+          if Map.has_key?(allowances, false), do: Map.delete(allowances, fun), else: allowances
 
-  defp resolve_once([], _kv, {result, resolved, unresolved}) do
-    {result, resolved, unresolved}
-  end
-
-  defp resolve_once(_, kv, {result, resolved, _unresolved}) do
-    {[kv | result], resolved, true}
-  end
-
-  defp fix_resolved({_, [], _}, state), do: state
-
-  defp fix_resolved({allowances, _fun_to_pids, lazy_calls}, state) do
-    allowances =
-      Enum.reduce(allowances, %{}, fn {k, v}, acc ->
-        Map.update(acc, k, v, &Map.merge(&1, v))
+        _, allowances ->
+          allowances
       end)
 
-    %__MODULE__{state | allowances: allowances, lazy_calls: lazy_calls}
+    %{state | allowances: updated_allowances}
   end
 end

test/nimble_ownership_test.exs

@@ -83,6 +83,31 @@ defmodule NimbleOwnershipTest do
       assert Exception.message(error) =~ "this PID is already allowed to access key"
     end
 
+    test "returns an error if PID A allows PID X (lazily) and  PID X tries to update",
+         %{key: key} do
+      owner_pid = spawn(fn -> Process.sleep(:infinity) end)
+
+      init_key(owner_pid, key, 1)
+
+      assert :ok =
+               NimbleOwnership.allow(
+                 @server,
+                 owner_pid,
+                 fn -> Process.whereis(:self_name) end,
+                 key
+               )
+
+      Process.register(self(), :self_name)
+
+      assert {:error, error} =
+               NimbleOwnership.get_and_update(@server, self(), key, fn _ ->
+                 {:yeah, %{}}
+               end)
+
+      assert error == %Error{reason: {:already_allowed, owner_pid}, key: key}
+      assert Exception.message(error) =~ "this PID is already allowed to access key"
+    end
+
     test "can set and update keys in shared mode", %{key: key} do
       NimbleOwnership.set_mode_to_shared(@server, self())
 
@@ -278,6 +303,7 @@ defmodule NimbleOwnershipTest do
               end)
 
               send(parent_pid, :done)
+              Process.sleep(:infinity)
           end
         end
 
@@ -289,13 +315,13 @@ defmodule NimbleOwnershipTest do
       send(lazy_pid_1, :go)
       assert_receive :done
 
-      assert NimbleOwnership.fetch_owner(@server, [self()], key) == {:ok, self()}
+      assert NimbleOwnership.fetch_owner(@server, [lazy_pid_1], key) == {:ok, self()}
       assert get_meta(self(), key) == %{counter: 2}
 
       send(lazy_pid_2, :go)
       assert_receive :done
 
-      assert NimbleOwnership.fetch_owner(@server, [self()], key) == {:ok, self()}
+      assert NimbleOwnership.fetch_owner(@server, [lazy_pid_2], key) == {:ok, self()}
       assert get_meta(self(), key) == %{counter: 3}
     end
 
@@ -340,7 +366,7 @@ defmodule NimbleOwnershipTest do
         for _ <- 1..2 do
           receive do
             {:go, parent_pid, key} ->
-              assert {:ok, owner_pid} = NimbleOwnership.fetch_owner(@server, [parent_pid], key)
+              assert {:ok, owner_pid} = NimbleOwnership.fetch_owner(@server, [self()], key)
               assert owner_pid == parent_pid
 
               NimbleOwnership.get_and_update(@server, owner_pid, key, fn info ->
@@ -381,6 +407,43 @@ defmodule NimbleOwnershipTest do
       assert NimbleOwnership.fetch_owner(@server, [owner_pid], key) == {:ok, owner_pid}
     end
 
+    test "does not resolve irrelevant lazy PIDs", %{key: key} do
+      parent_pid = self()
+
+      key2 = :"#{key}-2"
+
+      # Init the key.
+      init_key(parent_pid, key, %{counter: 1})
+      init_key(parent_pid, key2, %{})
+
+      # Allow a lazy PID that will resolve later.
+      assert :ok =
+               NimbleOwnership.allow(
+                 @server,
+                 self(),
+                 fn -> Process.whereis(:lazy_pid) end,
+                 key
+               )
+
+      # Allow a lazy PID that will never resolver
+      assert :ok =
+               NimbleOwnership.allow(
+                 @server,
+                 self(),
+                 fn -> raise "Too eager!" end,
+                 key2
+               )
+
+      {:ok, lazy_pid} =
+        Task.start_link(fn ->
+          Process.sleep(:infinity)
+        end)
+
+      Process.register(lazy_pid, :lazy_pid)
+
+      assert NimbleOwnership.fetch_owner(@server, [lazy_pid], key) == {:ok, self()}
+    end
+
     test "is idempotent", %{key: key} do
       owner_pid = spawn(fn -> Process.sleep(:infinity) end)