Use hazelnut sets instead of dictionaries to avoid duplicate hazels

This commit is contained in:
Yohann D'ANELLO 2021-01-05 20:26:22 +01:00
parent c7c8765fad
commit f062ee7b13
Signed by: ynerant
GPG Key ID: 3A75C55819C8CF85
2 changed files with 17 additions and 18 deletions

View File

@ -115,12 +115,12 @@ class Squirrel(Hazelnut):
return hazelnut return hazelnut
@property @property
def active_hazelnuts(self) -> dict: def active_hazelnuts(self) -> set:
return {(address, port): hazelnut for (address, port), hazelnut in self.hazelnuts.items() if hazelnut.active} return set(hazelnut for hazelnut in self.hazelnuts.values() if hazelnut.active)
@property @property
def potential_hazelnuts(self) -> dict: def potential_hazelnuts(self) -> set:
return {(address, port): hazelnut for (address, port), hazelnut in self.hazelnuts.items() if hazelnut.potential} return set(hazelnut for hazelnut in self.hazelnuts.values() if hazelnut.potential)
def find_hazelnut(self, address: str, port: int) -> Hazelnut: def find_hazelnut(self, address: str, port: int) -> Hazelnut:
""" """
@ -284,7 +284,7 @@ class Squirrel(Hazelnut):
self.add_message(msg) self.add_message(msg)
pkt = Packet.construct(DataTLV.construct(msg, self)) pkt = Packet.construct(DataTLV.construct(msg, self))
for hazelnut in list(self.active_hazelnuts.values()): for hazelnut in self.active_hazelnuts:
self.send_packet(hazelnut, pkt) self.send_packet(hazelnut, pkt)
def handle_mouse_click(self, y: int, x: int, attr: int) -> None: def handle_mouse_click(self, y: int, x: int, attr: int) -> None:
@ -369,7 +369,7 @@ class Squirrel(Hazelnut):
Takes the active hazels dictionnary and returns a list of [hazel, date+random, 0] Takes the active hazels dictionnary and returns a list of [hazel, date+random, 0]
""" """
res = dict() res = dict()
hazels = list(self.active_hazelnuts.items()) hazels = self.active_hazelnuts
for key, hazel in hazels: for key, hazel in hazels:
if hazel.symmetric: if hazel.symmetric:
next_send = uniform(1, 2) next_send = uniform(1, 2)
@ -622,8 +622,8 @@ class Squirrel(Hazelnut):
not have enough symmetric neighbours. not have enough symmetric neighbours.
""" """
res = [] res = []
lp = len(self.potential_hazelnuts) val = list(self.potential_hazelnuts)
val = list(self.potential_hazelnuts.values()) lp = len(val)
for i in range(min(lp, max(0, self.minNS - self.nbNS))): for i in range(min(lp, max(0, self.minNS - self.nbNS))):
a = randint(0, lp - 1) a = randint(0, lp - 1)
@ -634,7 +634,7 @@ class Squirrel(Hazelnut):
""" """
Sends a long HelloTLV to all active neighbours. Sends a long HelloTLV to all active neighbours.
""" """
for hazelnut in self.active_hazelnuts.values(): for hazelnut in self.active_hazelnuts:
htlv = HelloTLV().construct(16, self, hazelnut) htlv = HelloTLV().construct(16, self, hazelnut)
pkt = Packet().construct(htlv) pkt = Packet().construct(htlv)
self.send_packet(hazelnut, pkt) self.send_packet(hazelnut, pkt)
@ -644,7 +644,7 @@ class Squirrel(Hazelnut):
All neighbours that have not sent a HelloTLV in the last 2 All neighbours that have not sent a HelloTLV in the last 2
minutes are considered not active. minutes are considered not active.
""" """
val = list(self.active_hazelnuts.values()) # create a copy because the dict size will change val = list(self.active_hazelnuts) # create a copy because the dict size will change
for hazelnut in val: for hazelnut in val:
if time.time() - hazelnut.last_hello_time > 2 * 60: if time.time() - hazelnut.last_hello_time > 2 * 60:
@ -681,17 +681,17 @@ class Squirrel(Hazelnut):
nb_ns = 0 nb_ns = 0
# could send the same to all neighbour, but it means that neighbour # could send the same to all neighbour, but it means that neighbour
# A could receive a message with itself in it -> if the others do not pay attention, trouble # A could receive a message with itself in it -> if the others do not pay attention, trouble
for key, hazelnut in self.active_hazelnuts.items(): for hazelnut in self.active_hazelnuts:
if time.time() - hazelnut.last_long_hello_time <= 2 * 60: if time.time() - hazelnut.last_long_hello_time <= 2 * 60:
nb_ns += 1 nb_ns += 1
self.hazelnuts[key].symmetric = True hazelnut.symmetric = True
ntlv = NeighbourTLV().construct(*hazelnut.main_address) ntlv = NeighbourTLV().construct(*hazelnut.main_address)
pkt = Packet().construct(ntlv) pkt = Packet().construct(ntlv)
for destination in self.active_hazelnuts.values(): for destination in self.active_hazelnuts:
if destination.id != hazelnut.id: if destination.id != hazelnut.id:
self.send_packet(destination, pkt) self.send_packet(destination, pkt)
else: else:
self.active_hazelnuts[key].symmetric = False hazelnut.symmetric = False
self.nbNS = nb_ns self.nbNS = nb_ns
def leave(self) -> None: def leave(self) -> None:
@ -705,7 +705,7 @@ class Squirrel(Hazelnut):
# Broadcast a GoAway # Broadcast a GoAway
gatlv = GoAwayTLV().construct(GoAwayType.EXIT, "I am leaving! Good bye!") gatlv = GoAwayTLV().construct(GoAwayType.EXIT, "I am leaving! Good bye!")
pkt = Packet.construct(gatlv) pkt = Packet.construct(gatlv)
for hazelnut in self.active_hazelnuts.values(): for hazelnut in self.active_hazelnuts:
self.send_packet(hazelnut, pkt) self.send_packet(hazelnut, pkt)
exit(0) exit(0)
@ -762,7 +762,7 @@ class HazelManager(Thread):
# Second part: send long HelloTLVs to neighbours every 30 seconds # Second part: send long HelloTLVs to neighbours every 30 seconds
if time.time() - self.last_check > 30: if time.time() - self.last_check > 30:
self.squirrel.add_system_message(f"I have {len(list(self.squirrel.active_hazelnuts))} friends") self.squirrel.add_system_message(f"I have {len(self.squirrel.active_hazelnuts)} friends")
self.squirrel.send_hello() self.squirrel.send_hello()
self.last_check = time.time() self.last_check = time.time()

View File

@ -233,8 +233,7 @@ class NeighbourTLV(TLV):
# This case should never happen (and in our protocol it is not possible), # This case should never happen (and in our protocol it is not possible),
# but we include this test as a security measure. # but we include this test as a security measure.
return return
if not (str(self.ip_address), self.port) in squirrel.active_hazelnuts \ if not (str(self.ip_address), self.port) in squirrel.hazelnuts:
and not (str(self.ip_address), self.port) in squirrel.potential_hazelnuts:
hazelnut = squirrel.new_hazel(str(self.ip_address), self.port) hazelnut = squirrel.new_hazel(str(self.ip_address), self.port)
hazelnut.potential = True hazelnut.potential = True
squirrel.update_hazelnut_table(hazelnut) squirrel.update_hazelnut_table(hazelnut)