[syft/network] update NodePeer objects using a get and update with …

…`SyftLock` in network stash

packages/syft/src/syft/service/network/network_service.py

@@ -22,6 +22,8 @@
 from ...store.document_store import PartitionKey
 from ...store.document_store import PartitionSettings
 from ...store.document_store import QueryKeys
+from ...store.locks import NoLockingConfig
+from ...store.locks import SyftLock
 from ...types.grid_url import GridURL
 from ...types.transforms import TransformContext
 from ...types.transforms import keep
@@ -77,6 +79,7 @@ class NetworkStash(BaseUIDStoreStash):
 
  def __init__(self, store: DocumentStore) -> None:
  super().__init__(store=store)
+ self.lock = SyftLock(NoLockingConfig())
 
  def get_by_name(
  self, credentials: SyftVerifyKey, name: str
@@ -95,6 +98,25 @@ def update(
  return Err(SyftError(message=valid.err()))
  return super().update(credentials, peer, has_permission=has_permission)
 
+ def get_and_update_with_lock(
+ self,
+ credentials: SyftVerifyKey,
+ peer: NodePeer,
+ has_permission: bool = False,
+ ) -> Result[NodePeer, str]:
+ with self.lock:
+ result = self.get_by_uid(credentials=credentials, uid=peer.id)
+ if result.is_err():
+ return Err(
+ f"Failed to query peer {peer.id} with name '{peer.name}' from stash. Err: {result.err()}"
+ )
+ existing = result.ok()
+ if existing is None:
+ return Err(
+ f"Failed to query peer {peer.id} with name '{peer.name}' from stash: peer is None"
+ )
+ return self.update(credentials, peer, has_permission=has_permission)
+
  def create_or_update_peer(
  self, credentials: SyftVerifyKey, peer: NodePeer
  ) -> Result[NodePeer, str]:
@@ -136,7 +158,7 @@ def update_peer_ping_status(
  result = self.get_by_uid(credentials=credentials, uid=peer.id)
  if result.is_err():
  return Err(
- f"Failed to query peer peer {peer.id} with name '{peer.name}' from stash. Err: {result.err()}"
+ f"Failed to query peer {peer.id} with name '{peer.name}' from stash. Err: {result.err()}"
  )
  existing = result.ok()
  if existing is None:
@@ -612,7 +634,7 @@ def add_route(
  f"peer '{remote_node_peer.name}' with id '{existed_route.id}'."
  )
  # update the peer in the store with the updated routes
- result = self.stash.update(
+ result = self.stash.get_and_update_with_lock(
  credentials=context.node.verify_key,
  peer=remote_node_peer,
  )
@@ -770,7 +792,7 @@ def delete_route(
  )
  else:
  # update the peer with the route removed
- result = self.stash.update(
+ result = self.stash.get_and_update_with_lock(
  credentials=context.node.verify_key, peer=remote_node_peer
  )
  if result.is_err():
@@ -869,7 +891,9 @@ def update_route_priority(
  return updated_node_route
  new_priority: int = updated_node_route.priority
  # update the peer in the store
- result = self.stash.update(context.node.verify_key, remote_node_peer)
+ result = self.stash.get_and_update_with_lock(
+ context.node.verify_key, remote_node_peer
+ )
  if result.is_err():
  return SyftError(message=str(result.err()))
 

packages/syft/src/syft/service/network/utils.py

@@ -84,7 +84,7 @@ def peer_route_heathcheck(self, context: AuthedServiceContext) -> SyftError | No
  else:
  peer.ping_status_message = f"Peer '{peer.name}''s ping status: {peer.ping_status.value.lower()}"
 
- result = network_stash.update_peer_ping_status(
+ result = network_stash.get_and_update_with_lock(
  credentials=context.node.verify_key,
  peer=peer,
  has_permission=True,