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