SessMgr/Distr: improve handling of session hand-over

lib/floof.ex

@@ -106,7 +106,6 @@ defmodule Floof do
 
     case dcd do
       {:session, {:attach, key}} ->
-        Floof.Distributor.register_key(key, self())
         Floof.SessionManager.attach(key, self())
         Floof.SessionManager.set_multi(sesskey, backlog)
         backlog

lib/floof/distributor.ex

@@ -17,7 +17,7 @@ defmodule Floof.Distributor do
     pubkey_config: %{:keydb => ""}
   )
 
-  Record.defrecord(:dstate, fconf: {}, trgs: MapSet.new(), keys: %{}, seen: nil)
+  Record.defrecord(:dstate, fconf: {}, trgs: MapSet.new(), keys: MapSet.new(), seen: nil)
 
   def start_link(initval) do
     GenServer.start_link(__MODULE__, initval, name: __MODULE__)
@@ -51,32 +51,20 @@ defmodule Floof.Distributor do
 
   @impl true
   def handle_cast({:register, pid}, state) do
-    Process.link(pid)
-    {:noreply, dstate(state, trgs: MapSet.put(dstate(state, :trgs), pid))}
+    if :erlang.is_process_alive(pid) do
+      Process.link(pid)
+      {:noreply, dstate(state, trgs: MapSet.put(dstate(state, :trgs), pid))}
+    else
+      {:noreply, state}
+    end
   end
 
   @impl true
   def handle_cast({:registerKey, key, pid}, state) do
     trgs = dstate(state, :trgs)
     keys = dstate(state, :keys)
-
-    trgs =
-      case Map.get(keys, key) do
-        nil ->
-          trgs
-
-        oldpid ->
-          if :erlang.is_process_alive(oldpid) do
-            # the processes in trgs are always linked, those in keys aren't
-            Process.link(oldpid)
-            Process.unlink(pid)
-            MapSet.put(trgs, oldpid)
-          else
-            trgs
-          end
-      end
-
-    {:noreply, dstate(state, trgs: MapSet.delete(trgs, pid), keys: Map.put(keys, key, pid))}
+    Process.unlink(pid)
+    {:noreply, dstate(state, trgs: MapSet.delete(trgs, pid), keys: MapSet.put(keys, key))}
   end
 
   @impl true

lib/floof/session_manager.ex

@@ -6,8 +6,11 @@ defmodule Floof.SessionManager do
   end
 
   def attach(key, pid) do
+    Floof.Distributor.register_key(key, pid)
+
     modify_subscription(key, fn oldpid, m ->
       if oldpid != nil and oldpid != pid do
+        Floof.Distributor.register(oldpid)
         send(oldpid, {:SessionDetached, key})
       end