[neo3] Add Ping expiration (#1829)

* PingPong expiration * ms * Change to TimeProvider.Current.UtcNow * Optimize
neo-project · Aug 12, 2020 · ca13358 · ca13358
1 parent 6ded6df
commit ca13358
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 13 deletions.
diff --git a/src/neo/Network/P2P/RemoteNode.ProtocolHandler.cs b/src/neo/Network/P2P/RemoteNode.ProtocolHandler.cs
@@ -292,6 +292,7 @@ private void OnInventoryReceived(IInventory inventory)
             system.Blockchain.Tell(inventory, ActorRefs.NoSender);
             pendingKnownHashes.Remove(inventory.Hash);
             knownHashes.Add(inventory.Hash);
+            if (inventory is Block b) UpdateLastBlockIndex(b.Index, false);
         }
 
         private void OnInvMessageReceived(InvPayload payload)
@@ -323,13 +324,13 @@ private void OnMemPoolMessageReceived()
 
         private void OnPingMessageReceived(PingPayload payload)
         {
-            UpdateLastBlockIndex(payload);
+            UpdateLastBlockIndex(payload.LastBlockIndex, true);
             EnqueueMessage(Message.Create(MessageCommand.Pong, PingPayload.Create(Blockchain.Singleton.Height, payload.Nonce)));
         }
 
         private void OnPongMessageReceived(PingPayload payload)
         {
-            UpdateLastBlockIndex(payload);
+            UpdateLastBlockIndex(payload.LastBlockIndex, true);
         }
 
         private void OnVerackMessageReceived()
@@ -375,12 +376,12 @@ private void RefreshPendingKnownHashes()
             }
         }
 
-        private void UpdateLastBlockIndex(PingPayload payload)
+        private void UpdateLastBlockIndex(uint lastBlockIndex, bool requestTasks)
         {
-            if (payload.LastBlockIndex > LastBlockIndex)
+            if (lastBlockIndex > LastBlockIndex)
             {
-                LastBlockIndex = payload.LastBlockIndex;
-                system.TaskManager.Tell(new TaskManager.Update { LastBlockIndex = LastBlockIndex });
+                LastBlockIndex = lastBlockIndex;
+                system.TaskManager.Tell(new TaskManager.Update { LastBlockIndex = LastBlockIndex, RequestTasks = requestTasks });
             }
         }
     }

diff --git a/src/neo/Network/P2P/TaskManager.cs b/src/neo/Network/P2P/TaskManager.cs
@@ -15,7 +15,7 @@ namespace Neo.Network.P2P
     internal class TaskManager : UntypedActor
     {
         public class Register { public VersionPayload Version; }
-        public class Update { public uint LastBlockIndex; }
+        public class Update { public uint LastBlockIndex; public bool RequestTasks; }
         public class NewTasks { public InvPayload Payload; }
         public class RestartTasks { public InvPayload Payload; }
         private class Timer { }
@@ -129,7 +129,7 @@ protected override void OnReceive(object message)
                     OnRegister(register.Version);
                     break;
                 case Update update:
-                    OnUpdate(update.LastBlockIndex);
+                    OnUpdate(update);
                     break;
                 case NewTasks tasks:
                     OnNewTasks(tasks.Payload);
@@ -169,12 +169,13 @@ private void OnRegister(VersionPayload version)
             RequestTasks();
         }
 
-        private void OnUpdate(uint lastBlockIndex)
+        private void OnUpdate(Update update)
         {
             if (!sessions.TryGetValue(Sender, out TaskSession session))
                 return;
-            session.LastBlockIndex = lastBlockIndex;
-            RequestTasks();
+            session.LastBlockIndex = update.LastBlockIndex;
+            session.ExpireTime = TimeProvider.Current.UtcNow.AddMilliseconds(PingCoolingOffPeriod);
+            if (update.RequestTasks) RequestTasks();
         }
 
         private void OnRestartTasks(InvPayload payload)
@@ -302,14 +303,17 @@ private void SendPingMessage()
             {
                 var node = item.Key;
                 var session = item.Value;
-                if (Blockchain.Singleton.Height >= session.LastBlockIndex
-                    && TimeProvider.Current.UtcNow.ToTimestampMS() - PingCoolingOffPeriod >= Blockchain.Singleton.GetBlock(Blockchain.Singleton.CurrentBlockHash)?.Timestamp)
+
+                if (session.ExpireTime < TimeProvider.Current.UtcNow ||
+                     (Blockchain.Singleton.Height >= session.LastBlockIndex
+                     && TimeProvider.Current.UtcNow.ToTimestampMS() - PingCoolingOffPeriod >= Blockchain.Singleton.GetBlock(Blockchain.Singleton.CurrentBlockHash)?.Timestamp))
                 {
                     if (session.InvTasks.Remove(MemPoolTaskHash))
                     {
                         node.Tell(Message.Create(MessageCommand.Mempool));
                     }
                     node.Tell(Message.Create(MessageCommand.Ping, PingPayload.Create(Blockchain.Singleton.Height)));
+                    session.ExpireTime = TimeProvider.Current.UtcNow.AddMilliseconds(PingCoolingOffPeriod);
                 }
             }
         }

diff --git a/src/neo/Network/P2P/TaskSession.cs b/src/neo/Network/P2P/TaskSession.cs
@@ -15,6 +15,7 @@ internal class TaskSession
         public uint LastBlockIndex { get; set; }
         public uint TimeoutTimes = 0;
         public uint InvalidBlockCount = 0;
+        public DateTime ExpireTime = DateTime.MinValue;
 
         public TaskSession(VersionPayload version)
         {