Fix akka warning (#1533)

neo-project · Apr 10, 2020 · cb34f79 · cb34f79
1 parent 8b558d6
commit cb34f79
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 15 deletions.
diff --git a/src/neo/Network/P2P/LocalNode.cs b/src/neo/Network/P2P/LocalNode.cs
@@ -1,6 +1,5 @@
 using Akka.Actor;
 using Neo.IO;
-using Neo.Ledger;
 using Neo.Network.P2P.Payloads;
 using System;
 using System.Collections.Concurrent;
@@ -214,6 +213,11 @@ private void OnRelayDirectly(IInventory inventory)
 
         private void OnSendDirectly(IInventory inventory) => SendToRemoteNodes(inventory);
 
+        protected override void OnTcpConnected(IActorRef connection)
+        {
+            connection.Tell(new RemoteNode.StartProtocol());
+        }
+
         public static Props Props(NeoSystem system)
         {
             return Akka.Actor.Props.Create(() => new LocalNode(system));

diff --git a/src/neo/Network/P2P/Peer.cs b/src/neo/Network/P2P/Peer.cs
@@ -226,9 +226,14 @@ private void OnTcpConnected(IPEndPoint remote, IPEndPoint local)
                 Context.Watch(connection);
                 Sender.Tell(new Tcp.Register(connection));
                 ConnectedPeers.TryAdd(connection, remote);
+                OnTcpConnected(connection);
             }
         }
 
+        protected virtual void OnTcpConnected(IActorRef connection)
+        {
+        }
+
         /// <summary>
         /// Will be triggered when a Tcp.CommandFailed message is received.
         /// If it's a Tcp.Connect command, remove the related endpoint from ConnectingPeers.

diff --git a/src/neo/Network/P2P/RemoteNode.cs b/src/neo/Network/P2P/RemoteNode.cs
@@ -16,6 +16,7 @@ namespace Neo.Network.P2P
 {
     public partial class RemoteNode : Connection
     {
+        internal class StartProtocol { }
         internal class Relay { public IInventory Inventory; }
 
         private readonly NeoSystem system;
@@ -35,16 +36,6 @@ public RemoteNode(NeoSystem system, object connection, IPEndPoint remote, IPEndP
         {
             this.system = system;
             LocalNode.Singleton.RemoteNodes.TryAdd(Self, this);
-
-            var capabilities = new List<NodeCapability>
-            {
-                new FullNodeCapability(Blockchain.Singleton.Height)
-            };
-
-            if (LocalNode.Singleton.ListenerTcpPort > 0) capabilities.Add(new ServerCapability(NodeCapabilityType.TcpServer, (ushort)LocalNode.Singleton.ListenerTcpPort));
-            if (LocalNode.Singleton.ListenerWsPort > 0) capabilities.Add(new ServerCapability(NodeCapabilityType.WsServer, (ushort)LocalNode.Singleton.ListenerWsPort));
-
-            SendMessage(Message.Create(MessageCommand.Version, VersionPayload.Create(LocalNode.Nonce, LocalNode.UserAgent, capabilities.ToArray())));
         }
 
         /// <summary>
@@ -141,6 +132,9 @@ protected override void OnReceive(object message)
                 case Relay relay:
                     OnRelay(relay.Inventory);
                     break;
+                case StartProtocol _:
+                    OnStartProtocol();
+                    break;
             }
         }
 
@@ -166,6 +160,19 @@ private void OnSend(IInventory inventory)
             EnqueueMessage((MessageCommand)inventory.InventoryType, inventory);
         }
 
+        private void OnStartProtocol()
+        {
+            var capabilities = new List<NodeCapability>
+            {
+                new FullNodeCapability(Blockchain.Singleton.Height)
+            };
+
+            if (LocalNode.Singleton.ListenerTcpPort > 0) capabilities.Add(new ServerCapability(NodeCapabilityType.TcpServer, (ushort)LocalNode.Singleton.ListenerTcpPort));
+            if (LocalNode.Singleton.ListenerWsPort > 0) capabilities.Add(new ServerCapability(NodeCapabilityType.WsServer, (ushort)LocalNode.Singleton.ListenerWsPort));
+
+            SendMessage(Message.Create(MessageCommand.Version, VersionPayload.Create(LocalNode.Nonce, LocalNode.UserAgent, capabilities.ToArray())));
+        }
+
         protected override void PostStop()
         {
             timer.CancelIfNotNull();

diff --git a/tests/neo.UnitTests/Network/P2P/UT_RemoteNode.cs b/tests/neo.UnitTests/Network/P2P/UT_RemoteNode.cs
@@ -31,8 +31,6 @@ public void RemoteNode_Test_Abort_DifferentMagic()
             var connectionTestProbe = CreateTestProbe();
             var remoteNodeActor = ActorOfAsTestActorRef(() => new RemoteNode(testBlockchain, connectionTestProbe, null, null));
 
-            connectionTestProbe.ExpectMsg<Tcp.Write>();
-
             var msg = Message.Create(MessageCommand.Version, new VersionPayload
             {
                 UserAgent = "".PadLeft(1024, '0'),
@@ -58,8 +56,6 @@ public void RemoteNode_Test_Accept_IfSameMagic()
             var connectionTestProbe = CreateTestProbe();
             var remoteNodeActor = ActorOfAsTestActorRef(() => new RemoteNode(testBlockchain, connectionTestProbe, null, null));
 
-            connectionTestProbe.ExpectMsg<Tcp.Write>();
-
             var msg = Message.Create(MessageCommand.Version, new VersionPayload()
             {
                 UserAgent = "Unit Test".PadLeft(1024, '0'),