1 files changed, 7 insertions, 3 deletions
diff --git a/shard/lib/net/manager.ex b/shard/lib/net/manager.ex
index 75307ee..624d7e6 100644
--- a/shard/lib/net/manager.ex
+++ b/shard/lib/net/manager.ex
@@ -22,8 +22,11 @@ defmodule SNet.Manager do
     {:ok, nil} 
   end
 
-  def handle_call({:add_peer, peer_info, auth}, _from, state) do
+  def handle_call({:add_peer, peer_info, auth, callback}, _from, state) do
     pid = add_peer_internal(peer_info, auth)
+    if callback != nil do
+      GenServer.cast(pid, {:callback, callback})
+    end
     {:reply, pid, state}
   end
 
@@ -38,6 +41,7 @@ defmodule SNet.Manager do
       [[pid2]|_] when pid2 != pid ->
         {:reply, :redundant, state}
       _ ->
+        :ets.match_delete(:connections, {peer_info, pid, :_})
         :ets.insert(:connections, {peer_info, pid, auth})
 
         # Send interested message for all our shards
@@ -81,8 +85,8 @@ defmodule SNet.Manager do
   @doc"""
   Connect to a peer specified by ip address and port
   """
-  def add_peer(peer_info, auth \\ nil) do
-    GenServer.call(__MODULE__, {:add_peer, peer_info, auth})
+  def add_peer(peer_info, opts \\ []) do
+    GenServer.call(__MODULE__, {:add_peer, peer_info, opts[:auth], opts[:callback]})
   end
 
   @doc"""