RD-4775 - support threading in tracer (#212)

* support threading in tracer

* add threading tests

* avoid re-escaping

* fix Nirhod's CR

* fix Dori CR

Author: saartochner-lumigo

src/lumigo_tracer/lumigo_utils.py

@@ -67,6 +67,7 @@
 SKIP_SCRUBBING_KEYS = [EXECUTION_TAGS_KEY]
 LUMIGO_SECRET_MASKING_REGEX_BACKWARD_COMP = "LUMIGO_BLACKLIST_REGEX"
 LUMIGO_SECRET_MASKING_REGEX = "LUMIGO_SECRET_MASKING_REGEX"
+LUMIGO_SYNC_TRACING = "LUMIGO_SYNC_TRACING"
 WARN_CLIENT_PREFIX = "Lumigo Warning"
 INTERNAL_ANALYTICS_PREFIX = "Lumigo Analytic Log"
 TRUNCATE_SUFFIX = "...[too long]"
@@ -134,6 +135,7 @@ class Configuration:
     edge_kinesis_aws_access_key_id: Optional[str] = None
     edge_kinesis_aws_secret_access_key: Optional[str] = None
     should_scrub_known_services: bool = False
+    is_sync_tracer: bool = False
 
     @staticmethod
     def get_max_entry_size(has_error: bool = False) -> int:
@@ -239,6 +241,7 @@ def config(
     Configuration.should_scrub_known_services = (
         os.environ.get("LUMIGO_SCRUB_KNOWN_SERVICES") == "true"
     )
+    Configuration.is_sync_tracer = os.environ.get(LUMIGO_SYNC_TRACING, "FALSE").lower() == "true"
 
 
 def _is_span_has_error(span: dict) -> bool:
@@ -747,6 +750,8 @@ def lumigo_dumps(
         return "[" + ", ".join(organs) + "]"
 
     try:
+        if isinstance(d, str) and d.endswith(TRUNCATE_SUFFIX):
+            return d
         retval = aws_dump(d, decimal_safe=decimal_safe)
     except TypeError:
         if enforce_jsonify:
@@ -757,6 +762,22 @@ def lumigo_dumps(
     )
 
 
+def concat_old_body_to_new(old_body: Optional[str], new_body: bytes) -> str:
+    """
+    We have only a dumped body from the previous request,
+    so to concatenate the new body we should undo the lumigo_dumps.
+    Note that the old body is dumped bytes
+    """
+    if not new_body:
+        return old_body or ""
+    if not old_body:
+        return lumigo_dumps(new_body)
+    if old_body.endswith(TRUNCATE_SUFFIX):
+        return old_body
+    undumped_body = (old_body or "").encode().strip(b'"')
+    return lumigo_dumps(undumped_body + new_body)
+
+
 def is_kill_switch_on():
     return str(os.environ.get(KILL_SWITCH, "")).lower() == "true"
 

src/lumigo_tracer/spans_container.py

@@ -105,7 +105,7 @@ def __init__(
             self.base_msg,
         )
         self.span_ids_to_send: Set[str] = set()
-        self.spans: List[Dict] = []
+        self.spans: Dict[str, Dict] = {}
         if is_new_invocation:
             SpansContainer.is_cold = False
 
@@ -127,7 +127,7 @@ def start(self, event=None, context=None):
 
     def handle_timeout(self, *args):
         get_logger().info("The tracer reached the end of the timeout timer")
-        to_send = [s for s in self.spans if s["id"] in self.span_ids_to_send]
+        to_send = [self.spans[span_id] for span_id in self.span_ids_to_send]
         self.span_ids_to_send.clear()
         if Configuration.send_only_if_error:
             to_send.append(self._generate_start_span())
@@ -150,44 +150,48 @@ def add_span(self, span: dict) -> dict:
         This function parses an request event and add it to the span.
         """
         new_span = recursive_json_join(span, self.base_msg)
-        self.spans.append(new_span)
-        self.span_ids_to_send.add(span["id"])
+        span_id = new_span["id"]
+        self.spans[span_id] = new_span
+        self.span_ids_to_send.add(span_id)
         return new_span
 
-    def get_last_span(self) -> Optional[dict]:
-        if not self.spans:
+    def get_span_by_id(self, span_id: Optional[str]) -> Optional[dict]:
+        if not span_id:
             return None
-        return self.spans[-1]
+        return self.spans.get(span_id)
 
-    def get_span_by_id(self, span_id: str) -> Optional[dict]:
-        for span in self.spans:
-            if span.get("id") == span_id:
-                return span
-        return None
-
-    def pop_last_span(self) -> Optional[dict]:
-        return self.spans.pop() if self.spans else None
+    def pop_span(self, span_id: Optional[str]) -> Optional[dict]:
+        if not span_id:
+            return None
+        self.span_ids_to_send.discard(span_id)
+        return self.spans.pop(span_id, None)
 
-    def update_event_end_time(self) -> None:
+    def update_event_end_time(self, span_id: str) -> None:
         """
         This function assumes synchronous execution - we update the last http event.
         """
-        if self.spans:
-            span = self.spans[-1]
-            span["ended"] = get_current_ms_time()
-            self.span_ids_to_send.add(span["id"])
+        if span_id in self.spans:
+            self.spans[span_id]["ended"] = get_current_ms_time()
+            self.span_ids_to_send.add(span_id)
+        else:
+            get_logger().warning(f"update_event_end_time: Got unknown span id: {span_id}")
 
     def update_event_times(
-        self, start_time: Optional[datetime] = None, end_time: Optional[datetime] = None
+        self,
+        span_id: str,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
     ) -> None:
         """
         This function assumes synchronous execution - we update the last http event.
         """
-        if self.spans:
+        if span_id in self.spans:
             start_timestamp = start_time.timestamp() if start_time else time.time()
             end_timestamp = end_time.timestamp() if end_time else time.time()
-            self.spans[-1]["started"] = int(start_timestamp * 1000)
-            self.spans[-1]["ended"] = int(end_timestamp * 1000)
+            self.spans[span_id]["started"] = int(start_timestamp * 1000)
+            self.spans[span_id]["ended"] = int(end_timestamp * 1000)
+        else:
+            get_logger().warning(f"update_event_times: Got unknown span id: {span_id}")
 
     @staticmethod
     def _create_exception_event(
@@ -223,8 +227,9 @@ def add_exception_event(
     def add_step_end_event(self, ret_val):
         message_id = str(uuid.uuid4())
         step_function_span = create_step_function_span(message_id)
-        self.spans.append(recursive_json_join(step_function_span, self.base_msg))
-        self.span_ids_to_send.add(step_function_span["id"])
+        span_id = step_function_span["id"]
+        self.spans[span_id] = recursive_json_join(step_function_span, self.base_msg)
+        self.span_ids_to_send.add(span_id)
         if isinstance(ret_val, dict):
             ret_val[LUMIGO_EVENT_KEY] = {STEP_FUNCTION_UID_KEY: message_id}
             get_logger().debug(f"Added key {LUMIGO_EVENT_KEY} to the user's return value")
@@ -259,12 +264,12 @@ def end(self, ret_val=None, event: Optional[dict] = None, context=None) -> Optio
         if _is_span_has_error(self.function_span):
             self._set_error_extra_data(event)
         spans_contain_errors: bool = any(
-            _is_span_has_error(s) for s in self.spans + [self.function_span]
-        )
+            _is_span_has_error(s) for s in self.spans.values()
+        ) or _is_span_has_error(self.function_span)
 
         if (not Configuration.send_only_if_error) or spans_contain_errors:
             to_send = [self.function_span] + [
-                s for s in self.spans if s["id"] in self.span_ids_to_send
+                span for span_id, span in self.spans.items() if span_id in self.span_ids_to_send
             ]
             reported_rtt = lumigo_utils.report_json(region=self.region, msgs=to_send)
         else:

src/lumigo_tracer/wrappers/http/http_data_classes.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import Optional, List
+from typing import Optional, List, Dict
 
 
 class HttpRequest:
@@ -27,10 +27,13 @@ def clone(self, **kwargs):
 
 class HttpState:
     previous_request: Optional[HttpRequest] = None
-    previous_response_body: bytes = b""
+    previous_span_id: Optional[str] = None
     omit_skip_path: Optional[List[str]] = None
+    request_id_to_span_id: Dict[int, str] = {}
+    response_id_to_span_id: Dict[int, str] = {}
 
     @staticmethod
     def clear():
         HttpState.previous_request = None
-        HttpState.previous_response_body = b""
+        HttpState.request_id_to_span_id.clear()
+        HttpState.response_id_to_span_id.clear()