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

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,6 +21,7 @@ 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:
with lock:
edge = None edge = None
try: try:
edge = EdgeService.objects.get(src_node_id=src_node_id, dst_node_id=dst_node_id) edge = EdgeService.objects.get(src_node_id=src_node_id, dst_node_id=dst_node_id)