All local elements of PG work correctly

Author: michalmuskala

lib/firenest/pg.ex

@@ -5,29 +5,40 @@ defmodule Firenest.PG do
 
   defdelegate child_spec(opts), to: Firenest.PG.Supervisor
 
-  def track(pg, pid, group, key, meta) when node(pid) == node() do
+  def join(pg, group, key, pid, meta) when node(pid) == node() do
     server = partition_info!(pg, group)
-    GenServer.call(server, {:track, pid, group, key, meta})
+    GenServer.call(server, {:join, group, key, pid, meta})
   end
 
-  def untrack(pg, pid, group, key) when node(pid) == node() do
+  def leave(pg, group, key, pid) when node(pid) == node() do
     server = partition_info!(pg, group)
-    GenServer.call(server, {:untrack, pid, group, key})
+    GenServer.call(server, {:leave, group, key, pid})
   end
 
-  def untrack(pg, pid) when node(pid) == node() do
+  def leave(pg, pid) when node(pid) == node() do
     servers = partition_infos!(pg)
-    multicall(servers, {:untrack, pid}, 5_000)
+    replies = multicall(servers, {:leave, pid}, 5_000)
+
+    if :ok in replies do
+      :ok
+    else
+      {:error, :not_member}
+    end
   end
 
-  def update(pg, pid, group, key, meta) when node(pid) == node() do
+  def update(pg, group, key, pid, update) when node(pid) == node() and is_function(update, 1) do
     server = partition_info!(pg, group)
-    GenServer.call(server, {:update, pid, group, key, meta})
+    GenServer.call(server, {:update, group, key, pid, update})
   end
 
-  def list(pg, group) do
+  def replace(pg, group, key, pid, meta) when node(pid) == node() do
     server = partition_info!(pg, group)
-    GenServer.call(server, {:list, group}).()
+    GenServer.call(server, {:replace, group, key, pid, meta})
+  end
+
+  def members(pg, group) do
+    server = partition_info!(pg, group)
+    GenServer.call(server, {:members, group}).()
   end
 
   # TODO
@@ -41,7 +52,7 @@ defmodule Firenest.PG do
       send(pid, {:"$gen_call", {self(), ref}, request})
       ref
     end)
-    |> Enum.each(fn ref ->
+    |> Enum.map(fn ref ->
       receive do
         {^ref, reply} ->
           Process.demonitor(ref, [:flush])
@@ -134,22 +145,28 @@ defmodule Firenest.PG.Server do
   end
 
   @impl true
-  def handle_call({:track, pid, group, key, meta}, _from, state) do
+  def handle_call({:join, group, key, pid, meta}, _from, state) do
     %{values: values, pids: pids} = state
     Process.link(pid)
-    :ets.insert(values, {{group, pid, key}, meta})
-    :ets.insert(pids, {group, pid, key})
-    {:reply, :ok, state}
+    ets_key = {group, pid, key}
+
+    if :ets.member(values, ets_key) do
+      {:reply, {:error, :already_joined}, state}
+    else
+      :ets.insert(values, {{group, pid, key}, meta})
+      :ets.insert(pids, {group, pid, key})
+      {:reply, :ok, state}
+    end
   end
 
-  def handle_call({:untrack, pid, group, key}, _from, state) do
+  def handle_call({:leave, group, key, pid}, _from, state) do
     %{values: values, pids: pids} = state
     key = {group, pid, key}
     ms = [{key, [], [true]}]
 
     case :ets.select_delete(pids, ms) do
       0 ->
-        {:reply, {:error, :not_tracked}, state}
+        {:reply, {:error, :not_member}, state}
 
       1 ->
         unless :ets.member(pids, pid) do
@@ -161,18 +178,44 @@ defmodule Firenest.PG.Server do
     end
   end
 
-  def handle_call({:untrack, pid}, _from, state) do
+  def handle_call({:update, group, key, pid, update}, _from, state) do
+    %{values: values} = state
+    ets_key = {group, pid, key}
+
+    case ets_fetch_element(values, ets_key, 2) do
+      {:ok, value} ->
+        :ets.insert(values, {ets_key, update.(value)})
+        {:reply, :ok, state}
+
+      :error ->
+        {:reply, {:error, :not_member}, state}
+    end
+  end
+
+  def handle_call({:replace, group, key, pid, meta}, _from, state) do
+    %{values: values} = state
+    ets_key = {group, pid, key}
+
+    if :ets.member(values, ets_key) do
+      :ets.insert(values, {ets_key, meta})
+      {:reply, :ok, state}
+    else
+      {:reply, {:error, :not_member}, state}
+    end
+  end
+
+  def handle_call({:leave, pid}, _from, state) do
     %{values: values, pids: pids} = state
 
     if untrack_pid(pids, values, pid) do
       Process.unlink(pid)
       {:reply, :ok, state}
     else
-      {:reply, {:error, :not_tracked}, state}
+      {:reply, {:error, :not_member}, state}
     end
   end
 
-  def handle_call({:list, group}, _from, state) do
+  def handle_call({:members, group}, _from, state) do
     %{values: values} = state
 
     read = fn ->
@@ -205,4 +248,10 @@ defmodule Firenest.PG.Server do
         true
     end
   end
+
+  defp ets_fetch_element(table, key, pos) do
+    {:ok, :ets.lookup_element(table, key, pos)}
+  catch
+    :error, :badarg -> :error
+  end
 end

test/firenest/pg_test.exs

@@ -12,49 +12,125 @@ defmodule Firenest.PGTest do
     {:ok, pg: test}
   end
 
-  test "tracks processes", %{pg: pg} do
-    PG.track(pg, self(), :foo, :bar, :baz)
-    assert [{:bar, :baz}] == PG.list(pg, :foo)
+  describe "join/5" do
+    test "adds process", %{pg: pg} do
+      assert PG.join(pg, :foo, :bar, self(), :baz) == :ok
+      assert [{:bar, :baz}] == PG.members(pg, :foo)
+    end
+
+    test "rejects double joins", %{pg: pg} do
+      assert PG.join(pg, :foo, :bar, self(), :baz) == :ok
+      assert PG.join(pg, :foo, :bar, self(), :baz) == {:error, :already_joined}
+    end
+
+    test "cleans up entries after process dies", %{pg: pg} do
+      {pid, ref} = spawn_monitor(Process, :sleep, [:infinity])
+      PG.join(pg, :foo, :bar, pid, :baz)
+      assert [_] = PG.members(pg, :foo)
+      Process.exit(pid, :kill)
+      assert_receive {:DOWN, ^ref, _, _, _}
+      assert [] = PG.members(pg, :foo)
+    end
+
+    test "pg dies if other linked process dies", %{pg: pg} do
+      parent = self()
+      [{_, pid, _, _}] = Supervisor.which_children(Module.concat(pg, "Supervisor"))
+      ref = Process.monitor(pid)
+
+      temp =
+        spawn(fn ->
+          Process.link(pid)
+          send(parent, :continue)
+          Process.sleep(:infinity)
+        end)
+
+      assert_receive :continue
+
+      Process.exit(temp, :shutdown)
+      assert_receive {:DOWN, ^ref, _, _, _}
+    end
   end
 
-  test "processes are clened up when they die", %{pg: pg} do
-    {pid, ref} = spawn_monitor(Process, :sleep, [:infinity])
-    PG.track(pg, pid, :foo, :bar, :baz)
-    assert [_] = PG.list(pg, :foo)
-    Process.exit(pid, :kill)
-    assert_receive {:DOWN, ^ref, _, _, _}
-    assert [] = PG.list(pg, :foo)
+  describe "leave/2" do
+    test "removes entry", %{pg: pg} do
+      PG.join(pg, :foo, :bar, self(), :baz)
+
+      assert [_] = PG.members(pg, :foo)
+      assert PG.leave(pg, self()) == :ok
+      assert [] == PG.members(pg, :foo)
+    end
+
+    test "does not remove non members", %{pg: pg} do
+      [{_, pid, _, _}] = Supervisor.which_children(Module.concat(pg, "Supervisor"))
+      Process.link(pid)
+
+      assert PG.leave(pg, self()) == {:error, :not_member}
+      {:links, links} = Process.info(self(), :links)
+      assert pid in links
+    end
   end
 
-  test "pg dies if linked, untracked process terminates", %{pg: pg} do
-    parent = self()
-    [{_, pid, _, _}] = Supervisor.which_children(Module.concat(pg, "Supervisor"))
-    ref = Process.monitor(pid)
+  describe "leave/4" do
+    test "removes single entry", %{pg: pg} do
+      PG.join(pg, :foo, :bar, self(), :baz)
+      assert [_] = PG.members(pg, :foo)
+
+      assert PG.leave(pg, :foo, :bar, self()) == :ok
+      assert [] == PG.members(pg, :foo)
+    end
 
-    temp =
-      spawn(fn ->
-        Process.link(pid)
-        send(parent, :continue)
-        Process.sleep(:infinity)
-      end)
+    test "leaves other entries intact", %{pg: pg} do
+      PG.join(pg, :foo, :bar, self(), :baz)
+      PG.join(pg, :foo, :baar, self(), :baz)
+      assert [_, _] = PG.members(pg, :foo)
 
-    assert_receive :continue
+      assert PG.leave(pg, :foo, :bar, self()) == :ok
+      assert [{:baar, :baz}] == PG.members(pg, :foo)
+    end
 
-    Process.exit(temp, :shutdown)
-    assert_receive {:DOWN, ^ref, _, _, _}
+    test "does not remove non members", %{pg: pg} do
+      [{_, pid, _, _}] = Supervisor.which_children(Module.concat(pg, "Supervisor"))
+      Process.link(pid)
+
+      assert PG.leave(pg, :foo, :bar, self()) == {:error, :not_member}
+      {:links, links} = Process.info(self(), :links)
+      assert pid in links
+    end
   end
 
-  test "untrack/2", %{pg: pg} do
-    PG.track(pg, self(), :foo, :bar, :baz)
-    assert [_] = PG.list(pg, :foo)
-    PG.untrack(pg, self())
-    assert [] == PG.list(pg, :foo)
+  describe "update/5" do
+    test "executes the update if entry is present", %{pg: pg} do
+      parent = self()
+      PG.join(pg, :foo, :bar, self(), 1)
+      assert [{:bar, 1}] == PG.members(pg, :foo)
+
+      update = fn value ->
+        send(parent, value)
+        value + 1
+      end
+
+      assert PG.update(pg, :foo, :bar, self(), update) == :ok
+      assert_received 1
+      assert [{:bar, 2}] == PG.members(pg, :foo)
+    end
+
+    test "does not execute update if entry is absent", %{pg: pg} do
+      parent = self()
+      update = fn value -> Process.exit(parent, {:unexpected_update, value}) end
+      assert PG.update(pg, :foo, :bar, self(), update) == {:error, :not_member}
+    end
   end
 
-  test "untrack/4", %{pg: pg} do
-    PG.track(pg, self(), :foo, :bar, :baz)
-    assert [_] = PG.list(pg, :foo)
-    PG.untrack(pg, self(), :foo, :bar)
-    assert [] == PG.list(pg, :foo)
+  describe "replace/5" do
+    test "updates value if entry is present", %{pg: pg} do
+      PG.join(pg, :foo, :bar, self(), 1)
+      assert [{:bar, 1}] == PG.members(pg, :foo)
+      assert PG.replace(pg, :foo, :bar, self(), 2) == :ok
+      assert [{:bar, 2}] == PG.members(pg, :foo)
+    end
+
+    test "does not update value if entry is absent", %{pg: pg} do
+      assert PG.replace(pg, :foo, :bar, self(), 2) == {:error, :not_member}
+    end
   end
 end