RD-4424 - cold start indication in pre-invocation calls (#177)

* cold start indication in pre-invocation calls

Author: saartochner-lumigo

src/lumigo_tracer/spans_container.py

@@ -55,6 +55,7 @@ def __init__(
         trace_id_suffix: str = None,
         trigger_by: dict = None,
         max_finish_time: int = None,
+        is_new_invocation: bool = False,
         event: str = None,
         envs: str = None,
     ):
@@ -102,7 +103,8 @@ def __init__(
         )
         self.span_ids_to_send: Set[str] = set()
         self.spans: List[Dict] = []
-        SpansContainer.is_cold = False
+        if is_new_invocation:
+            SpansContainer.is_cold = False
 
     def _generate_start_span(self) -> dict:
         to_send = self.function_span.copy()
@@ -285,14 +287,14 @@ def get_span(cls) -> "SpansContainer":
         return cls.create_span()
 
     @classmethod
-    def create_span(cls, event=None, context=None, force=False) -> "SpansContainer":
+    def create_span(cls, event=None, context=None, is_new_invocation=False) -> "SpansContainer":
         """
         This function creates a span out of a given AWS context.
         The force flag delete any existing span-container (to handle with warm execution of lambdas).
         Note that if lambda will be executed directly (regular pythonic function call and not invoked),
             it will override the container.
         """
-        if cls._span and not force:
+        if cls._span and not is_new_invocation:
             return cls._span
         # copy the event to ensure that we will not change it
         event = copy.deepcopy(event)
@@ -319,6 +321,7 @@ def create_span(cls, event=None, context=None, force=False) -> "SpansContainer":
             account=safe_split_get(getattr(context, "invoked_function_arn", ""), ":", 4, ""),
             trigger_by=parse_triggered_by(event),
             max_finish_time=get_current_ms_time() + remaining_time,
+            is_new_invocation=is_new_invocation,
             **additional_info,
         )
         return cls._span

src/lumigo_tracer/tracer.py

@@ -54,7 +54,7 @@ def lambda_wrapper(*args, **kwargs):
         try:
             if Configuration.enhanced_print:
                 _enhance_output(args, local_print, local_logging_format)
-            SpansContainer.create_span(*args, force=True)
+            SpansContainer.create_span(*args, is_new_invocation=True)
             with lumigo_safe_execute("auto tag"):
                 AutoTagEvent.auto_tag_event(args[0])
             SpansContainer.get_span().start(*args)

src/test/unit/test_tracer.py

@@ -463,3 +463,18 @@ def wrapped(event, context):
     from_lumigo = line_dropper.sub("-", stacktrace)
     original = line_dropper.sub("-", traceback.format_tb(e.value.__traceback__)[1])
     assert from_lumigo == original
+
+
+def test_cold_indicator_with_request_in_cold_phase(context):
+    SpansContainer.is_cold = True
+    #  Create a request the might invert the `is_cold` field
+    http.client.HTTPConnection("www.google.com").request("POST", "/")
+    assert SpansContainer.is_cold is True
+
+    @lumigo_tracer(step_function=True)
+    def lambda_test_function(event, context):
+        http.client.HTTPConnection("www.google.com").request("POST", "/")
+
+    lambda_test_function({}, context)
+    assert SpansContainer.get_span().function_span["readiness"] == "cold"
+    assert SpansContainer.is_cold is False