feat(Swarm): PeerDisconnected event

richardschneider · Apr 25, 2019 · 97d0fbc · 97d0fbc
1 parent 4537bee
commit 97d0fbc
Show file tree

Hide file tree

Showing 4 changed files with 107 additions and 0 deletions.
diff --git a/src/ConnectionManager.cs b/src/ConnectionManager.cs
@@ -31,6 +31,11 @@ public class ConnectionManager
         string Key(Peer peer) => peer.Id.ToBase58();
         string Key(MultiHash id) => id.ToBase58();
 
+        /// <summary>
+        ///   Raised when a peer's connection is closed.
+        /// </summary>
+        public event EventHandler<MultiHash> PeerDisconnected;
+
         /// <summary>
         ///   Gets the current connections.
         /// </summary>
@@ -165,6 +170,10 @@ public bool Remove(PeerConnection connection)
                 var last = newConns.Last();
                 last.RemotePeer.ConnectedAddress = last.RemoteAddress;
             }
+            else
+            {
+                PeerDisconnected?.Invoke(this, connection.RemotePeer.Id);
+            }
             return true;
         }
 
@@ -187,6 +196,8 @@ public bool Remove(MultiHash id)
             {
                 conn.Dispose();
             }
+
+            PeerDisconnected?.Invoke(this, id);
             return true;
         }
 

diff --git a/src/Swarm.cs b/src/Swarm.cs
@@ -64,6 +64,11 @@ public class Swarm : IService, IPolicy<MultiAddress>
         /// </summary>
         public event EventHandler<Peer> PeerDiscovered;
 
+        /// <summary>
+        ///   Raised when a peer's connection is closed.
+        /// </summary>
+        public event EventHandler<Peer> PeerDisconnected;
+
         /// <summary>
         ///  The local peer.
         /// </summary>
@@ -329,12 +334,22 @@ public Task StartAsync()
                 log.Warn("Peer key is missing, using unencrypted connections.");
             }
 
+            Manager.PeerDisconnected += OnPeerDisconnected;
             IsRunning = true;
             log.Debug("Started");
 
             return Task.CompletedTask;
         }
 
+        void OnPeerDisconnected(object sender, MultiHash peerId)
+        {
+            if (!otherPeers.TryGetValue(peerId.ToBase58(), out Peer peer))
+            {
+                peer = new Peer { Id = peerId };
+            }
+            PeerDisconnected?.Invoke(this, peer);
+        }
+
         /// <inheritdoc />
         public async Task StopAsync()
         {
@@ -349,6 +364,7 @@ public async Task StopAsync()
 
             // Disconnect from remote peers.
             Manager.Clear();
+            Manager.PeerDisconnected -= OnPeerDisconnected;
 
             otherPeers.Clear();
             listeners.Clear();

diff --git a/test/ConnectionManagerTest.cs b/test/ConnectionManagerTest.cs
@@ -265,5 +265,46 @@ public void Clear()
             Assert.IsNull(a.Stream);
             Assert.IsNull(b.Stream);
         }
+
+        [TestMethod]
+        public void PeerDisconnectedEvent_RemovingPeer()
+        {
+            bool gotEvent = false;
+            var manager = new ConnectionManager();
+            manager.PeerDisconnected += (s, e) => gotEvent = true;
+            var peerA = new Peer { Id = aId };
+            var a = new PeerConnection { RemotePeer = peerA, Stream = Stream.Null };
+            manager.Add(a);
+
+            manager.Remove(peerA.Id);
+            Assert.IsTrue(gotEvent);
+        }
+
+        [TestMethod]
+        public void PeerDisconnectedEvent_RemovingConnection()
+        {
+            int gotEvent = 0;
+            var manager = new ConnectionManager();
+            manager.PeerDisconnected += (s, e) => gotEvent += 1;
+            var peerA = new Peer { Id = aId };
+            var a = new PeerConnection { RemotePeer = peerA, Stream = Stream.Null };
+            manager.Add(a);
+
+            manager.Remove(a);
+            Assert.AreEqual(1, gotEvent);
+        }
+
+        [TestMethod]
+        public void PeerDisconnectedEvent_ConnectionClose()
+        {
+            int gotEvent = 0;
+            var manager = new ConnectionManager();
+            manager.PeerDisconnected += (s, e) => gotEvent += 1;
+            var peerA = new Peer { Id = aId };
+            var a = new PeerConnection { RemotePeer = peerA, Stream = Stream.Null };
+            manager.Add(a);
+            a.Dispose();
+            Assert.AreEqual(1, gotEvent);
+        }
     }
 }
diff --git a/test/SwarmTest.cs b/test/SwarmTest.cs
@@ -364,6 +364,45 @@ public void Connecting_To_Self()
             });
         }
 
+        [TestMethod]
+        public async Task PeerDisconnected()
+        {
+            var peerB = new Peer
+            {
+                AgentVersion = "peerB",
+                Id = "QmdpwjdB94eNm2Lcvp9JqoCxswo3AKQqjLuNZyLixmCM1h",
+                PublicKey = "CAASXjBcMA0GCSqGSIb3DQEBAQUAA0sAMEgCQQDlTSgVLprWaXfmxDr92DJE1FP0wOexhulPqXSTsNh5ot6j+UiuMgwb0shSPKzLx9AuTolCGhnwpTBYHVhFoBErAgMBAAE="
+            };
+            var swarmB = new Swarm { LocalPeer = peerB };
+            await swarmB.StartAsync();
+            var peerBAddress = await swarmB.StartListeningAsync("/ip4/127.0.0.1/tcp/0");
+
+            var swarm = new Swarm { LocalPeer = self };
+            var swarmConnections = 0;
+            swarm.ConnectionEstablished += (s, e) =>
+            {
+                ++swarmConnections;
+            };
+            swarm.PeerDisconnected += (s, e) =>
+            {
+                --swarmConnections;
+            };
+            await swarm.StartAsync();
+            try
+            {
+                var remotePeer = await swarm.ConnectAsync(peerBAddress);
+                Assert.AreEqual(1, swarmConnections);
+
+                await swarm.StopAsync();
+                Assert.AreEqual(0, swarmConnections);
+            }
+            finally
+            {
+                await swarm.StopAsync();
+                await swarmB.StopAsync();
+            }
+        }
+
         [TestMethod]
         public async Task Listening()
         {