p34709852
/
monkey
forked from p15670423/monkey
1
0
Fork 1
Code Issues Pull Requests Projects Releases Wiki Activity

Island: Use lock when creating an edge in the db 

Fixes #1917
PR #1932
This commit is contained in:
Shreya Malviya 2022-05-09 15:58:11 +05:30 committed by GitHub
parent 174392848f
commit 28f60d51ff
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 14 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
monkey/monkey_island/cc/services/edge/edge.py
View File

@ -1,6 +1,7 @@
from __future__ import annotations
import copy
import threading
from typing import Dict, List
from bson import ObjectId
@ -10,6 +11,8 @@ from monkey_island.cc.models.edge import Edge
RIGHT_ARROW = "\u2192"
lock = threading.Lock()
class EdgeService(Edge):
@staticmethod
@ -18,16 +21,17 @@ class EdgeService(Edge):
@staticmethod
def get_or_create_edge(src_node_id, dst_node_id, src_label, dst_label) -> EdgeService:
edge = None
try:
edge = EdgeService.objects.get(src_node_id=src_node_id, dst_node_id=dst_node_id)
except DoesNotExist:
edge = EdgeService(src_node_id=src_node_id, dst_node_id=dst_node_id)
finally:
if edge:
edge.update_label(node_id=src_node_id, label=src_label)
edge.update_label(node_id=dst_node_id, label=dst_label)
return edge
with lock:
edge = None
try:
edge = EdgeService.objects.get(src_node_id=src_node_id, dst_node_id=dst_node_id)
except DoesNotExist:
edge = EdgeService(src_node_id=src_node_id, dst_node_id=dst_node_id)
finally:
if edge:
edge.update_label(node_id=src_node_id, label=src_label)
edge.update_label(node_id=dst_node_id, label=dst_label)
return edge
@staticmethod
def get_by_dst_node(dst_node_id: ObjectId) -> List[EdgeService]: