fix(aws_lambda): fix cold start detection in provisioned concurrency functions

epsagon/wrappers/aws_lambda.py

@@ -3,6 +3,7 @@
 """
 
 from __future__ import absolute_import
+import os
 import traceback
 import time
 import copy
@@ -35,6 +36,7 @@ def _add_status_code(runner, return_value):
             runner.resource['metadata']['status_code'] = status_code
 
 
+# pylint: disable=too-many-statements
 def lambda_wrapper(func):
     """Epsagon's Lambda wrapper."""
 
@@ -58,6 +60,11 @@ def _lambda_wrapper(*args, **kwargs):
             # parameters / sends kwargs. In such case we ignore this trace.
             return func(*args, **kwargs)
 
+        if os.environ.get(
+            'AWS_LAMBDA_INITIALIZATION_TYPE'
+        ) == 'provisioned-concurrency':
+            constants.COLD_START = False
+
         try:
             runner = epsagon.runners.aws_lambda.LambdaRunner(
                 time.time(),

tests/wrappers/test_lambda_wrapper.py

@@ -1,3 +1,4 @@
+import os
 import json
 import mock
 import pytest
@@ -81,9 +82,53 @@ def wrapped_lambda(event, context):
     trace_mock.add_exception.assert_not_called()
 
     assert not epsagon.constants.COLD_START
+    assert runner.resource['metadata']['cold_start'] == True
     assert runner.resource['metadata']['return_value'] == retval
 
 
+@mock.patch.object(LambdaRunner, 'set_exception')
+@mock.patch(
+    'epsagon.trace.trace_factory.get_trace',
+    side_effect=lambda: trace_mock
+)
+@mock.patch(
+    'epsagon.trace.trace_factory.get_or_create_trace',
+    side_effect=lambda: trace_mock
+)
+@mock.patch(
+    'epsagon.triggers.aws_lambda.LambdaTriggerFactory.factory',
+    side_effect=['trigger']
+)
+def test_lambda_wrapper_provisioned_concurrency_sanity(
+        trigger_factory_mock,
+        _,
+        __,
+        set_exception_mock
+):
+    retval = 'success'
+
+    @epsagon.wrappers.aws_lambda.lambda_wrapper
+    def wrapped_lambda(event, context):
+        return 'success'
+
+    os.environ['AWS_LAMBDA_INITIALIZATION_TYPE'] = 'provisioned-concurrency'
+    assert wrapped_lambda('a', CONTEXT_STUB) == 'success'
+    trace_mock.prepare.assert_called()
+    runner = _get_runner_event(trace_mock)
+
+    trigger_factory_mock.assert_called()
+    set_exception_mock.assert_not_called()
+
+    trace_mock.set_timeout_handler.assert_called()
+
+    trace_mock.send_traces.assert_called()
+    trace_mock.add_exception.assert_not_called()
+
+    assert not epsagon.constants.COLD_START
+    assert runner.resource['metadata']['cold_start'] == False
+    os.environ.pop('AWS_LAMBDA_INITIALIZATION_TYPE')
+
+
 @mock.patch(
     'epsagon.trace.trace_factory.get_trace',
     side_effect=lambda: trace_mock