[neo2] Raise priority of Ping/Pong (#1812)

neo-project · Aug 10, 2020 · 95596f9 · 95596f9
1 parent 0a9f0db
commit 95596f9
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/neo/Network/P2P/TaskManager.cs b/neo/Network/P2P/TaskManager.cs
@@ -7,6 +7,7 @@
 using Neo.Network.P2P.Payloads;
 using System;
 using System.Collections;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 using System.Runtime.CompilerServices;
@@ -40,6 +41,7 @@ private class Timer { }
 
         private readonly Dictionary<UInt256, int> globalTasks = new Dictionary<UInt256, int>();
         private readonly Dictionary<IActorRef, TaskSession> sessions = new Dictionary<IActorRef, TaskSession>();
+        private readonly ConcurrentDictionary<ActorPath, DateTime> _expiredTimes = new ConcurrentDictionary<ActorPath, DateTime>();
         private readonly ICancelable timer = Context.System.Scheduler.ScheduleTellRepeatedlyCancelable(TimerInterval, TimerInterval, Context.Self, new Timer(), ActorRefs.NoSender);
 
         private readonly UInt256 HeaderTaskHash = UInt256.Zero;
@@ -196,6 +198,7 @@ private void OnTerminated(IActorRef actor)
             if (!sessions.TryGetValue(actor, out TaskSession session))
                 return;
             sessions.Remove(actor);
+            _expiredTimes.TryRemove(session.RemoteNode.Path, out _);
             foreach (UInt256 hash in session.Tasks.Keys)
                 DecrementGlobalTask(hash);
         }
@@ -226,6 +229,11 @@ public static Props Props(NeoSystem system)
 
         private void RequestTasks(TaskSession session)
         {
+            if (!_expiredTimes.TryGetValue(session.RemoteNode.Path, out var expireTime) || expireTime < DateTime.UtcNow)
+            {
+                session.RemoteNode.Tell(Message.Create("ping", PingPayload.Create(Blockchain.Singleton.Height)));
+                _expiredTimes[session.RemoteNode.Path] = DateTime.UtcNow.AddSeconds(PingCoolingOffPeriod);
+            }
             if (session.HasTask) return;
             if (session.AvailableTasks.Count > 0)
             {
@@ -267,11 +275,6 @@ private void RequestTasks(TaskSession session)
                 }
                 session.RemoteNode.Tell(Message.Create("getblocks", GetBlocksPayload.Create(hash)));
             }
-            else if (Blockchain.Singleton.HeaderHeight >= session.LastBlockIndex
-                    && TimeProvider.Current.UtcNow.ToTimestamp() - PingCoolingOffPeriod >= Blockchain.Singleton.GetBlock(Blockchain.Singleton.CurrentHeaderHash)?.Timestamp)
-            {
-                session.RemoteNode.Tell(Message.Create("ping", PingPayload.Create(Blockchain.Singleton.Height)));
-            }
             if (!HasStateRootTask)
             {
                 var state_height = Math.Max(Blockchain.Singleton.StateHeight, (long)ProtocolSettings.Default.StateRootEnableIndex - 1);