From 6a01f46dc6d20408e8654e6ce96165fee4c7e405 Mon Sep 17 00:00:00 2001
From: Tim Cooper <tim@youtooz.com>
Date: Wed, 25 Mar 2026 15:45:09 +1000
Subject: [PATCH] Fix owner_unallow using wrong ref when switching owners
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When ownership_allow is called with unallow_existing: true to move a
process from one owner to another, owner_unallow was removing the
process from the caller's owner ref rather than the process's current
owner ref. This left the process in the old owner's allowed list.

When the old owner later exits, owner_down removes all processes in
its allowed list from checkouts — including the process that was
already re-allowed on a different owner. This causes an
OwnershipError on the next query because the process is no longer in
checkouts.

The fix looks up the process's current checkout entry to find its
actual ref and proxy, then cleans up the correct owner's allowed list.

Also fixes the ETS delete in owner_unallow which was passing
{pid, proxy} as the key instead of just pid — a no-op on a :set
table keyed by the first element.
---
 integration_test/ownership/manager_test.exs | 56 +++++++++++++++++++++
 lib/db_connection/ownership/manager.ex      | 44 +++++++++-------
 2 files changed, 81 insertions(+), 19 deletions(-)

diff --git a/integration_test/ownership/manager_test.exs b/integration_test/ownership/manager_test.exs
index 962cfa2..0ac3956 100644
--- a/integration_test/ownership/manager_test.exs
+++ b/integration_test/ownership/manager_test.exs
@@ -697,6 +697,62 @@ defmodule ManagerTest do
     end
   end
 
+  test "unallow_existing cleans up old owner so its exit doesn't revoke the new allowance" do
+    stack = List.duplicate({:ok, :state}, 2) ++ List.duplicate({:idle, :state}, 20)
+    {:ok, agent} = A.start_link(stack)
+    opts = [agent: agent, parent: self(), ownership_mode: :manual, pool_size: 2]
+    {:ok, pool} = P.start_link(opts)
+    parent = self()
+
+    # A long-lived worker process that gets re-allowed across owners.
+    worker =
+      spawn(fn ->
+        receive do
+          :done -> :ok
+        end
+      end)
+
+    # owner_a checks out and allows the worker
+    owner_a =
+      async_no_callers(fn ->
+        :ok = Ownership.ownership_checkout(pool, [])
+        :ok = Ownership.ownership_allow(pool, self(), worker, [])
+        send(parent, {:owner_ready, :a})
+        assert_receive :exit
+      end)
+
+    assert_receive {:owner_ready, :a}
+
+    # owner_b checks out and re-allows the worker via unallow_existing
+    owner_b =
+      async_no_callers(fn ->
+        :ok = Ownership.ownership_checkout(pool, [])
+        :ok = Ownership.ownership_allow(pool, self(), worker, unallow_existing: true)
+        send(parent, {:owner_ready, :b})
+        assert_receive :exit
+      end)
+
+    assert_receive {:owner_ready, :b}
+
+    # Verify worker has access via owner_b
+    assert P.run(pool, fn _ -> :ok end, [caller: worker] ++ opts)
+
+    # Kill owner_a. The worker was moved to owner_b via unallow_existing,
+    # so owner_a's exit must not revoke the worker's access.
+    send(owner_a.pid, :exit)
+    ref = Process.monitor(owner_a.pid)
+    assert_receive {:DOWN, ^ref, _, _, _}
+
+    # Let the pool manager process the :DOWN
+    Process.sleep(50)
+
+    # Worker should still have access through owner_b
+    assert P.run(pool, fn _ -> :ok end, [caller: worker] ++ opts)
+
+    send(owner_b.pid, :exit)
+    send(worker, :done)
+  end
+
   test "doesn't require a label to produce an ownership error" do
     {:ok, pool, opts} = start_pool()
 
diff --git a/lib/db_connection/ownership/manager.ex b/lib/db_connection/ownership/manager.ex
index 8de2678..a6f3103 100644
--- a/lib/db_connection/ownership/manager.ex
+++ b/lib/db_connection/ownership/manager.ex
@@ -331,28 +331,34 @@ defmodule DBConnection.Ownership.Manager do
     state
   end
 
-  defp owner_unallow(%{ets: ets, log: log} = state, caller, unallow, ref, proxy) do
-    if log do
-      Logger.log(log, fn ->
-        [
-          Util.inspect_pid(unallow),
-          " was unallowed by ",
-          Util.inspect_pid(caller),
-          " on proxy ",
-          Util.inspect_pid(proxy)
-        ]
-      end)
-    end
+  defp owner_unallow(%{ets: ets, log: log} = state, caller, unallow, _ref, _proxy) do
+    case Map.get(state.checkouts, unallow, :not_found) do
+      {_status, old_ref, old_proxy} ->
+        if log do
+          Logger.log(log, fn ->
+            [
+              Util.inspect_pid(unallow),
+              " was unallowed by ",
+              Util.inspect_pid(caller),
+              " on proxy ",
+              Util.inspect_pid(old_proxy)
+            ]
+          end)
+        end
 
-    state = update_in(state.checkouts, &Map.delete(&1, unallow))
+        state = update_in(state.checkouts, &Map.delete(&1, unallow))
 
-    state =
-      update_in(state.owners[ref], fn {proxy, caller, allowed} ->
-        {proxy, caller, List.delete(allowed, unallow)}
-      end)
+        state =
+          update_in(state.owners[old_ref], fn {proxy, caller, allowed} ->
+            {proxy, caller, List.delete(allowed, unallow)}
+          end)
 
-    ets && :ets.delete(ets, {unallow, proxy})
-    state
+        ets && :ets.delete(ets, unallow)
+        state
+
+      :not_found ->
+        state
+    end
   end
 
   defp owner_down(%{ets: ets, log: log} = state, ref) do