fix(framework) Fix minor issues in InMemoryLinkState (#4764)

adap · Jan 4, 2025 · c3f318e · c3f318e
1 parent 8308d6e
commit c3f318e
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/py/flwr/server/superlink/linkstate/in_memory_linkstate.py b/src/py/flwr/server/superlink/linkstate/in_memory_linkstate.py
@@ -440,7 +440,7 @@ def clear_supernode_auth_keys_and_credentials(self) -> None:
     def store_node_public_keys(self, public_keys: set[bytes]) -> None:
         """Store a set of `node_public_keys` in the link state."""
         with self.lock:
-            self.node_public_keys = public_keys
+            self.node_public_keys.update(public_keys)
 
     def store_node_public_key(self, public_key: bytes) -> None:
         """Store a `node_public_key` in the link state."""
@@ -449,7 +449,8 @@ def store_node_public_key(self, public_key: bytes) -> None:
 
     def get_node_public_keys(self) -> set[bytes]:
         """Retrieve all currently stored `node_public_keys` as a set."""
-        return self.node_public_keys
+        with self.lock:
+            return self.node_public_keys.copy()
 
     def get_run_ids(self) -> set[int]:
         """Retrieve all run IDs."""