Merge pull request #4 from chenquanyu/add-relayreason

add expire milliseconds for dictionary
bettybao1209 · Apr 1, 2020 · 0211568 · 0211568
2 parents def1ebf + ca3eadc
commit 0211568
Showing 1 changed file with 14 additions and 4 deletions.
diff --git a/src/RpcServer/RelayActor.cs b/src/RpcServer/RelayActor.cs
@@ -1,6 +1,7 @@
 using Akka.Actor;
 using Neo.Network.P2P.Payloads;
 using System.Collections.Generic;
+using System.Threading.Tasks;
 using static Neo.Ledger.Blockchain;
 
 namespace Neo.Plugins
@@ -9,10 +10,12 @@ public class RelayActor : UntypedActor
     {
         private readonly NeoSystem neoSystem;
         private readonly Dictionary<UInt256, IActorRef> senders = new Dictionary<UInt256, IActorRef>();
+        private readonly int expireMs;
 
-        public RelayActor(NeoSystem neoSystem)
+        public RelayActor(NeoSystem neoSystem, int expireMs)
         {
             this.neoSystem = neoSystem;
+            this.expireMs = expireMs;
             Context.System.EventStream.Subscribe(Self, typeof(RelayResult));
         }
 
@@ -22,8 +25,15 @@ protected override void OnReceive(object message)
             {
                 case IInventory inventory:
                     {
-                        senders.Add(inventory.Hash, Sender);
+                        UInt256 hash = inventory.Hash;
+                        senders.Add(hash, Sender);
                         neoSystem.Blockchain.Tell(inventory);
+                        // the sender will clean the record after expireMs
+                        Task.Run(async () =>
+                        {
+                            await Task.Delay(expireMs);
+                            senders.Remove(hash);
+                        });
                         break;
                     }
                 case RelayResult reason:
@@ -37,9 +47,9 @@ protected override void OnReceive(object message)
             }
         }
 
-        public static Props Props(NeoSystem neoSystem)
+        public static Props Props(NeoSystem neoSystem, int expire = 10000)
         {
-            return Akka.Actor.Props.Create(() => new RelayActor(neoSystem));
+            return Akka.Actor.Props.Create(() => new RelayActor(neoSystem, expire));
         }
     }
 }