From 6c9f23fd8cd428e4db8f0f86d27d9e71cf17e6b2 Mon Sep 17 00:00:00 2001
From: saartochner <saart@lumigo.io>
Date: Sun, 27 Dec 2020 12:17:43 +0200
Subject: [PATCH] use initialization type to determine readiness

---
 src/lumigo_tracer/lumigo_utils.py    |  4 ++++
 src/lumigo_tracer/spans_container.py |  4 +++-
 src/test/unit/test_tracer.py         | 13 +++++++++++++
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/lumigo_tracer/lumigo_utils.py b/src/lumigo_tracer/lumigo_utils.py
index bc977a07..ba90303f 100644
--- a/src/lumigo_tracer/lumigo_utils.py
+++ b/src/lumigo_tracer/lumigo_utils.py
@@ -661,3 +661,7 @@ def is_error_code(status_code: int) -> bool:
 
 def is_aws_arn(string_to_validate: Optional[str]) -> bool:
     return bool(string_to_validate and string_to_validate.startswith("arn:aws:"))
+
+
+def is_provision_concurrency_initialization() -> bool:
+    return os.environ.get("AWS_LAMBDA_INITIALIZATION_TYPE") == "provisioned-concurrency"
diff --git a/src/lumigo_tracer/spans_container.py b/src/lumigo_tracer/spans_container.py
index 2572871a..1e587376 100644
--- a/src/lumigo_tracer/spans_container.py
+++ b/src/lumigo_tracer/spans_container.py
@@ -23,6 +23,7 @@
     create_step_function_span,
     get_current_ms_time,
     get_region,
+    is_provision_concurrency_initialization,
 )
 from lumigo_tracer import lumigo_utils
 from lumigo_tracer.parsing_utils import parse_trace_id, safe_split_get, recursive_json_join
@@ -78,6 +79,7 @@ def __init__(
             "info": {"tracer": {"version": version}, "traceId": {"Root": trace_root}},
             "token": Configuration.token,
         }
+        is_cold = SpansContainer.is_cold and not is_provision_concurrency_initialization()
         self.function_span = recursive_json_join(
             {
                 "id": request_id,
@@ -87,7 +89,7 @@ def __init__(
                 "event": event,
                 "envs": envs,
                 "memoryAllocated": memory_allocated,
-                "readiness": "cold" if SpansContainer.is_cold else "warm",
+                "readiness": "cold" if is_cold else "warm",
                 "info": {
                     "logStreamName": log_stream_name,
                     "logGroupName": log_group_name,
diff --git a/src/test/unit/test_tracer.py b/src/test/unit/test_tracer.py
index bdef594c..6ff540b8 100644
--- a/src/test/unit/test_tracer.py
+++ b/src/test/unit/test_tracer.py
@@ -96,6 +96,19 @@ def lambda_test_function(event, context):
     assert span["return_value"] == '{"a": [1.0]}'
 
 
+def test_lambda_wrapper_provision_concurrency_is_warm(context, monkeypatch):
+    monkeypatch.setattr(SpansContainer, "is_cold", True)
+    monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "provisioned-concurrency")
+
+    @lumigo_tracer(token="123")
+    def lambda_test_function(event, context):
+        return {"a": "b"}
+
+    lambda_test_function({}, context)
+    span = SpansContainer.get_span().function_span
+    assert span["readiness"] == "warm"
+
+
 def test_kill_switch(monkeypatch, context):
     monkeypatch.setattr(os, "environ", {"LUMIGO_SWITCH_OFF": "true"})