diff --git a/lib/floof/session_manager.ex b/lib/floof/session_manager.ex
index 0ec701c..a3f72c2 100644
--- a/lib/floof/session_manager.ex
+++ b/lib/floof/session_manager.ex
@@ -6,36 +6,36 @@ defmodule Floof.SessionManager do
   end
 
   def attach(key, pid) do
-    if Agent.get_and_update(__MODULE__, fn state ->
-         Map.get_and_update(state, key, fn q ->
-           {oldpid, q2} = entry_dfl(q)
+    modify_subscription(key, fn oldpid, m ->
+      if oldpid != nil and oldpid != pid do
+        send(oldpid, {:SessionDetached, key})
+      end
 
-           if oldpid != nil and oldpid != pid do
-             send(oldpid, {:SessionDetached, key})
-           end
+      if not Enum.empty?(m) do
+        send(pid, {:SessionPushed, key})
+      end
 
-           {Enum.empty?(q2), {pid, q2}}
-         end)
-       end) do
-      send(pid, {:SessionPushed, key})
-    end
+      pid
+    end)
   end
 
   def detach(key, pid) do
+    modify_subscription(key, fn oldpid, _ ->
+      if oldpid == pid do
+        send(pid, {:SessionDetached, key})
+        nil
+      else
+        oldpid
+      end
+    end)
+  end
+
+  defp modify_subscription(key, handler) do
     Agent.cast(__MODULE__, fn state ->
       {_, state2} =
-        Map.get_and_update(state, key, fn q ->
-          {sub, q2} = entry_dfl(q)
-
-          sub2 =
-            if sub == pid do
-              send(pid, {:SessionDetached, key})
-              nil
-            else
-              sub
-            end
-
-          {nil, {sub2, q2}}
+        Map.get_and_update(state, key, fn ent ->
+          {oldpid, m} = entry_dfl(ent)
+          {nil, {handler.(oldpid, m), m}}
         end)
 
       state2
@@ -74,15 +74,15 @@ defmodule Floof.SessionManager do
   end
 
   def peek(key) do
-    Agent.get_and_update(__MODULE__, fn state ->
-      retval =
-        case Map.get(state, key) do
+    Agent.get(
+      __MODULE__,
+      &MapSet.new(
+        case Map.get(&1, key) do
           nil -> []
           {_, m} -> Map.keys(m)
         end
-
-      {MapSet.new(retval), state}
-    end)
+      )
+    )
   end
 
   def get(key, subkey) do