Store a single hazelnut table

squinnondation/hazel.py

@@ -24,6 +24,8 @@ class Hazelnut:
         self.last_hello_time = 0
         self.last_long_hello_time = 0
         self.symmetric = False
+        self.active = False
+        self.potential = False
 
         try:
             # Resolve DNS as an IPv6
@@ -77,8 +79,7 @@ class Squirrel(Hazelnut):
             curses.init_pair(i + 1, i, curses.COLOR_BLACK)
 
         # dictionnaries of neighbours
-        self.potentialhazelnuts = dict()
-        self.activehazelnuts = dict()
+        self.hazelnuts = dict()
         self.nbNS = 0
         self.minNS = 3  # minimal number of symmetric neighbours a squirrel needs to have.
 
@@ -96,22 +97,21 @@ class Squirrel(Hazelnut):
         hazelnut = Hazelnut(address=address, port=port)
         return hazelnut
 
-    def is_active(self, hazel: Hazelnut) -> bool:
-        return (hazel.address, hazel.port) in self.activehazelnuts
+    @property
+    def activehazelnuts(self) -> dict:
+        return {(address, port): hazelnut for (address, port), hazelnut in self.hazelnuts.items() if hazelnut.active}
 
-    def is_potential(self, hazel: Hazelnut) -> bool:
-        return (hazel.address, hazel.port) in self.potentialhazelnuts
-
-    def remove_from_potential(self, address: str, port: int) -> None:
-        self.potentialhazelnuts.pop((address, port), None)
+    @property
+    def potentialhazelnuts(self) -> dict:
+        return {(address, port): hazelnut for (address, port), hazelnut in self.hazelnuts.items() if hazelnut.potential}
 
     def find_hazelnut(self, address: str, port: int) -> Hazelnut:
         """
         Translate an address into a hazelnut. If this hazelnut does not exist,
         creates a new hazelnut.
         """
-        if (address, port) in self.activehazelnuts:
-            return self.activehazelnuts[(address, port)]
+        if (address, port) in self.hazelnuts:
+            return self.hazelnuts[(address, port)]
         hazelnut = Hazelnut(address=address, port=port)
         return hazelnut
 
@@ -402,9 +402,8 @@ class Squirrel(Hazelnut):
                 if self.recent_messages[key][2][key2][2] >= 5:  # the neighbour is not reactive enough
                     gatlv = GoAwayTLV().construct(GoAwayType.TIMEOUT, f"{self.id} No acknowledge")
                     pkt = Packet().construct(gatlv)
-                    self.send_packet(self.recent_messages[key][2][key2][0], pkt)
-                    self.activehazelnuts.pop(key2)
-                    self.potentialhazelnuts[key2] = self.recent_messages[key][2][key2][0]
+                    hazelnut = self.recent_messages[key][2][key2][0]
+                    self.send_packet(hazelnut, pkt)
                     self.recent_messages[key][2].pop(key2)
 
         self.refresh_lock.release()
@@ -648,11 +647,15 @@ class Squirrel(Hazelnut):
                 gatlv = GoAwayTLV().construct(GoAwayType.TIMEOUT, "you did not talk to me")
                 pkt = Packet().construct(gatlv)
                 self.send_packet(hazelnut, pkt)
-                self.activehazelnuts.pop((hazelnut.address, hazelnut.port))
-                self.potentialhazelnuts[(hazelnut.address, hazelnut.port)] = hazelnut
+                hazelnut.active = False
+                hazelnut.potential = True
+                self.update_hazelnut_table(hazelnut)
 
         self.refresh_lock.release()
 
+    def update_hazelnut_table(self, hazelnut: Hazelnut) -> None:
+        self.hazelnuts[(hazelnut.address, hazelnut.port)] = hazelnut
+
     def send_neighbours(self) -> None:
         """
         Update the number of symmetric neighbours and
@@ -666,7 +669,7 @@ class Squirrel(Hazelnut):
         for key, hazelnut in self.activehazelnuts.items():
             if time.time() - hazelnut.last_long_hello_time <= 2 * 60:
                 nb_ns += 1
-                self.activehazelnuts[key].symmetric = True
+                self.hazelnuts[key].symmetric = True
                 ntlv = NeighbourTLV().construct(hazelnut.address, hazelnut.port)
                 pkt = Packet().construct(ntlv)
                 for destination in self.activehazelnuts.values():
@@ -748,7 +751,7 @@ class HazelManager(Thread):
 
             # Second part: send long HelloTLVs to neighbours every 30 seconds
             if time.time() - self.last_check > 30:
-                self.squirrel.add_system_message(f"I have {len(list(self.squirrel.activehazelnuts.values()))} friends")
+                self.squirrel.add_system_message(f"I have {len(list(self.squirrel.activehazelnuts))} friends")
                 self.squirrel.send_hello()
                 self.last_check = time.time()
 

squinnondation/messages.py

@@ -159,7 +159,7 @@ class HelloTLV(TLV):
 
     def handle(self, squirrel: Any, sender: Any) -> None:
         time_h = time.time()
-        if not squirrel.is_active(sender):
+        if not sender.active:
             sender.id = self.source_id  # The sender we are given misses an id
             time_hl = time.time()
         else:
@@ -167,14 +167,13 @@ class HelloTLV(TLV):
         if self.is_long and self.dest_id == squirrel.id:
             time_hl = time.time()
 
-        # Make sure the sender is not in the potential hazelnuts
-        squirrel.remove_from_potential(sender.address, sender.port)
-
         # Add entry to/actualize the active hazelnuts dictionnary
         sender.last_hello_time = time_h
         sender.last_long_hello_time = time_hl
         sender.symmetric = True
-        squirrel.activehazelnuts[(sender.address, sender.port)] = sender
+        sender.potential = False
+        sender.active = True
+        squirrel.update_hazelnut_table(sender)
         squirrel.nbNS += 1
         squirrel.add_system_message(f"{self.source_id} sent me a Hello " + ("long" if self.is_long else "short"))
 
@@ -374,9 +373,10 @@ class GoAwayTLV(TLV):
             self.message.encode("UTF-8")[:self.length - 1]
 
     def handle(self, squirrel: Any, sender: Any) -> None:
-        if squirrel.is_active(sender):
-            squirrel.activehazelnuts.pop((sender.address, sender.port))
-            squirrel.potentialhazelnuts[(sender.address, sender.port)] = sender
+        if sender.active:
+            sender.active = False
+            sender.potential = True
+            squirrel.update_hazelnut_table(sender)
         squirrel.add_system_message("Some told me that he went away : " + self.message)
 
     @staticmethod