Feature: send execution tags on timeout (#258)

* new span type: enrichment_span (send execution_tags on timeouted lambda)

Co-authored-by: CircleCI <[email protected]>

Author: saartochner-lumigo

.secrets.baseline

@@ -207,7 +207,7 @@
         "filename": "src/test/unit/test_tracer.py",
         "hashed_secret": "dabe695c9892b5cb937379817679c0fb0be6a3df",
         "is_verified": false,
-        "line_number": 422
+        "line_number": 427
       }
     ],
     "src/test/unit/wrappers/http/test_http_parser.py": [
@@ -220,5 +220,5 @@
       }
     ]
   },
-  "generated_at": "2022-08-22T12:01:52Z"
+  "generated_at": "2022-08-25T09:52:57Z"
 }

src/lumigo_tracer/spans_container.py

@@ -6,7 +6,7 @@
 import time
 import uuid
 import signal
-from typing import List, Dict, Optional, Callable, Set
+from typing import List, Dict, Optional, Callable, Set, Union
 
 from lumigo_tracer.event.event_dumper import EventDumper
 from lumigo_tracer.lumigo_utils import (
@@ -36,6 +36,7 @@
 _VERSION_PATH = os.path.join(os.path.dirname(__file__), "VERSION")
 MAX_LAMBDA_TIME = 15 * 60 * 1000
 FUNCTION_TYPE = "function"
+ENRICHMENT_TYPE = "enrichment"
 MALFORMED_TXID = "000000000000000000000000"
 
 
@@ -103,11 +104,17 @@ def __init__(
                     **(trigger_by or {}),
                 },
                 "isMalformedTransactionId": malformed_txid,
-                EXECUTION_TAGS_KEY: [],
                 MANUAL_TRACES_KEY: [],
             },
             self.base_msg,
         )
+        self.base_enrichment_span = {
+            "type": ENRICHMENT_TYPE,
+            "token": Configuration.token,
+            "invocation_id": request_id,
+            "transaction_id": transaction_id,
+        }
+        self.execution_tags: List[Dict[str, str]] = []
         self.span_ids_to_send: Set[str] = set()
         self.spans: Dict[str, Dict] = {}
         self.manual_trace_start_times: Dict[str, int] = {}
@@ -121,6 +128,14 @@ def _generate_start_span(self) -> dict:
         to_send["maxFinishTime"] = self.max_finish_time
         return to_send
 
+    def generate_enrichment_span(self) -> Optional[Dict[str, Union[str, int]]]:
+        if not self.execution_tags:
+            return None
+        return recursive_json_join(
+            {"sending_time": get_current_ms_time(), EXECUTION_TAGS_KEY: self.execution_tags.copy()},
+            self.base_enrichment_span,
+        )
+
     def start(self, event=None, context=None):
         to_send = self._generate_start_span()
         if not Configuration.send_only_if_error:
@@ -137,6 +152,9 @@ def handle_timeout(self, *args):
             get_logger().info("The tracer reached the end of the timeout timer")
             spans_id_copy = self.span_ids_to_send.copy()
             to_send = [self.spans[span_id] for span_id in spans_id_copy]
+            enrichment_span = self.generate_enrichment_span()
+            if enrichment_span:
+                to_send.insert(0, enrichment_span)
             self.span_ids_to_send.clear()
             if Configuration.send_only_if_error:
                 to_send.append(self._generate_start_span())
@@ -244,10 +262,10 @@ def add_step_end_event(self, ret_val):
             get_logger().debug(f"Added key {LUMIGO_EVENT_KEY} to the user's return value")
 
     def get_tags_len(self) -> int:
-        return len(self.function_span[EXECUTION_TAGS_KEY])
+        return len(self.execution_tags)
 
     def add_tag(self, key: str, value: str) -> None:
-        self.function_span[EXECUTION_TAGS_KEY].append({"key": key, "value": value})
+        self.execution_tags.append({"key": key, "value": value})
 
     def start_manual_trace(self, name: str) -> None:
         now = get_current_ms_time()
@@ -292,6 +310,9 @@ def end(self, ret_val=None, event: Optional[dict] = None, context=None) -> Optio
             to_send = [self.function_span] + [
                 span for span_id, span in self.spans.items() if span_id in self.span_ids_to_send
             ]
+            enrichment_span = self.generate_enrichment_span()
+            if enrichment_span:
+                to_send.append(enrichment_span)
             reported_rtt = lumigo_utils.report_json(region=self.region, msgs=to_send)
         else:
             get_logger().debug(

src/test/unit/auto_tag/test_auto_tag_event.py

@@ -9,7 +9,7 @@
     ConfigurationHandler,
 )
 from lumigo_tracer.spans_container import SpansContainer
-from lumigo_tracer.lumigo_utils import EXECUTION_TAGS_KEY, Configuration
+from lumigo_tracer.lumigo_utils import Configuration
 
 
 class ExceptionHandler(EventAutoTagHandler):
@@ -25,15 +25,15 @@ def auto_tag(event):
 def test_auto_tag_event_is_none():
     AutoTagEvent.auto_tag_event(event=None)
 
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == []
+    assert SpansContainer.get_span().execution_tags == []
 
 
 def test_auto_tag_exception():
     event = {"a": 1}
 
     AutoTagEvent.auto_tag_event(event=event, handlers=[ExceptionHandler()])
 
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == []
+    assert SpansContainer.get_span().execution_tags == []
 
 
 def test_auto_tag_key_not_in_header(monkeypatch):
@@ -156,7 +156,7 @@ def test_auto_tag_key_not_in_header(monkeypatch):
 
     AutoTagEvent.auto_tag_event(event=event)
 
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == []
+    assert SpansContainer.get_span().execution_tags == []
 
 
 def test_auto_tag_key_in_header(monkeypatch):
@@ -279,7 +279,7 @@ def test_auto_tag_key_in_header(monkeypatch):
 
     AutoTagEvent.auto_tag_event(event=event)
 
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == [
+    assert SpansContainer.get_span().execution_tags == [
         {"key": "Accept", "value": "application/json, text/plain, */*"}
     ]
 
@@ -326,7 +326,7 @@ def test_configuration_handler_is_supported(config, event, expected):
 def test_configuration_handler_auto_tag(auto_tag_keys, event, result_tags):
     Configuration.auto_tag = auto_tag_keys
     ConfigurationHandler.auto_tag(event)
-    tags = SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY]
+    tags = SpansContainer.get_span().execution_tags
     assert len(tags) == len(result_tags)
     for tag in result_tags:
         assert tag in tags
@@ -335,7 +335,7 @@ def test_configuration_handler_auto_tag(auto_tag_keys, event, result_tags):
 def test_configuration_handler_auto_tag_failure(capsys):
     Configuration.auto_tag = [None, "key2"]
     ConfigurationHandler.auto_tag({"key1": datetime, "key2": "value"})
-    tags = SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY]
+    tags = SpansContainer.get_span().execution_tags
     assert tags == [{"key": "key2", "value": "value"}]
     assert "Failed to auto tag" in capsys.readouterr().out
 
@@ -352,5 +352,5 @@ def test_configuration_handler_auto_tag_non_string(value, expected):
 
     ConfigurationHandler.auto_tag({"key1": value})
 
-    tags = SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY]
+    tags = SpansContainer.get_span().execution_tags
     assert {"key": "key1", "value": expected} in tags

src/test/unit/test_spans_container.py

@@ -9,19 +9,20 @@
 
 import pytest
 
-from lumigo_tracer import lumigo_utils
+from lumigo_tracer import lumigo_utils, add_execution_tag
 from lumigo_tracer.wrappers.http.http_parser import HTTP_TYPE
 from lumigo_tracer.spans_container import (
     SpansContainer,
     TimeoutMechanism,
     FUNCTION_TYPE,
     MALFORMED_TXID,
+    ENRICHMENT_TYPE,
 )
 from lumigo_tracer.lumigo_utils import (
     Configuration,
-    EXECUTION_TAGS_KEY,
     MANUAL_TRACES_KEY,
     get_current_ms_time,
+    EXECUTION_TAGS_KEY,
 )
 
 
@@ -200,13 +201,35 @@ def test_timeout_mechanism_timeout_occurred_send_new_spans(monkeypatch, context,
     assert SpansContainer.get_span().span_ids_to_send
 
 
+def test_timeout_mechanism_timeout_occurred_but_finish_check_enrichment(
+    monkeypatch, context, dummy_span, reporter_mock
+):
+    SpansContainer.create_span()
+    SpansContainer.get_span().start(context=context)
+    SpansContainer.get_span().add_span(dummy_span)
+    add_execution_tag("key", "value")
+    SpansContainer.get_span().handle_timeout()
+
+    add_execution_tag("new_key", "new_value")
+    SpansContainer.get_span().end(ret_val={"hello": "world"})
+
+    first_send = reporter_mock.call_args_list[1][1]["msgs"]
+    enrichment_span = next(s for s in first_send if s["type"] == ENRICHMENT_TYPE)
+    assert enrichment_span[EXECUTION_TAGS_KEY] == [{"key": "key", "value": "value"}]
+
+    final_send = reporter_mock.call_args_list[-1][1]["msgs"]
+    enrichment_span = next(s for s in final_send if s["type"] == ENRICHMENT_TYPE)
+    assert enrichment_span[EXECUTION_TAGS_KEY] == [
+        {"key": "key", "value": "value"},
+        {"key": "new_key", "value": "new_value"},
+    ]
+
+
 def test_add_tag():
     key = "my_key"
     value = "my_value"
     SpansContainer.get_span().add_tag(key, value)
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == [
-        {"key": key, "value": value}
-    ]
+    assert SpansContainer.get_span().execution_tags == [{"key": key, "value": value}]
 
 
 def test_start_manual_trace_simple_flow():

src/test/unit/test_tracer.py

@@ -25,7 +25,7 @@
     SKIP_COLLECTING_HTTP_BODY_KEY,
 )
 
-from lumigo_tracer.spans_container import SpansContainer
+from lumigo_tracer.spans_container import SpansContainer, ENRICHMENT_TYPE
 from moto import mock_kinesis
 
 TOKEN = "t_10faa5e13e7844aaa1234"
@@ -365,7 +365,12 @@ def lambda_test_function(event, context):
     assert reporter_mock.call_count == 2
 
 
-def test_wrapping_with_tags(context):
+def get_enrichment_spans(reporter_mock):
+    final_send = reporter_mock.call_args_list[-1][1]["msgs"]
+    return [s for s in final_send if s["type"] == ENRICHMENT_TYPE]
+
+
+def test_wrapping_with_tags(context, reporter_mock):
     key = "my_key"
     value = "my_value"
 
@@ -376,25 +381,25 @@ def lambda_test_function(event, context):
 
     result = lambda_test_function({}, context)
     assert result == "ret_value"
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == [
-        {"key": key, "value": value}
-    ]
+    enrichment_spans = get_enrichment_spans(reporter_mock)
+    assert len(enrichment_spans) == 1
+    assert enrichment_spans[0][EXECUTION_TAGS_KEY] == [{"key": key, "value": value}]
 
 
 @pytest.mark.parametrize(
     "key, event",
     [("my_key", {"my_key": "my_value"}), ("my_key.key2", {"my_key": {"key2": "my_value"}})],
 )
-def test_wrapping_with_auto_tags(context, key, event):
+def test_wrapping_with_auto_tags(context, key, event, reporter_mock):
     @lumigo_tracer(auto_tag=[key])
     def lambda_test_function(event, context):
         return "ret_value"
 
     result = lambda_test_function(event, context)
     assert result == "ret_value"
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == [
-        {"key": key, "value": "my_value"}
-    ]
+    enrichment_spans = get_enrichment_spans(reporter_mock)
+    assert len(enrichment_spans) == 1
+    assert enrichment_spans[0][EXECUTION_TAGS_KEY] == [{"key": key, "value": "my_value"}]
 
 
 def test_not_jsonable_return(monkeypatch, context):

src/test/unit/test_user_utils.py

@@ -16,7 +16,7 @@
     manual_trace,
     manual_trace_sync,
 )
-from lumigo_tracer.lumigo_utils import EXECUTION_TAGS_KEY, MANUAL_TRACES_KEY
+from lumigo_tracer.lumigo_utils import MANUAL_TRACES_KEY
 
 
 def test_manual_traces_context_manager():
@@ -107,9 +107,7 @@ def test_add_execution_tag():
     key = "my_key"
     value = "my_value"
     assert add_execution_tag(key, value) is True
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == [
-        {"key": key, "value": value}
-    ]
+    assert SpansContainer.get_span().execution_tags == [{"key": key, "value": value}]
 
 
 def test_start_manual_trace_simple_flow():
@@ -121,25 +119,25 @@ def test_start_manual_trace_simple_flow():
 def test_add_execution_key_tag_empty(capsys):
     assert add_execution_tag("", "value") is False
     assert "Unable to add tag: key length" in capsys.readouterr().out
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == []
+    assert SpansContainer.get_span().execution_tags == []
 
 
 def test_add_execution_value_tag_empty(capsys):
     assert add_execution_tag("key", "") is False
     assert "Unable to add tag: value length" in capsys.readouterr().out
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == []
+    assert SpansContainer.get_span().execution_tags == []
 
 
 def test_add_execution_tag_key_pass_max_chars(capsys):
     assert add_execution_tag("k" * (MAX_TAG_KEY_LEN + 1), "value") is False
     assert "Unable to add tag: key length" in capsys.readouterr().out
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == []
+    assert SpansContainer.get_span().execution_tags == []
 
 
 def test_add_execution_tag_value_pass_max_chars(capsys):
     assert add_execution_tag("key", "v" * (MAX_TAG_VALUE_LEN + 1)) is False
     assert "Unable to add tag: value length" in capsys.readouterr().out
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == []
+    assert SpansContainer.get_span().execution_tags == []
 
 
 def test_add_execution_tag_pass_max_tags(capsys):
@@ -155,7 +153,7 @@ def test_add_execution_tag_pass_max_tags(capsys):
 
     assert "Unable to add tag: maximum number of tags" in capsys.readouterr().out
     assert (
-        SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY]
+        SpansContainer.get_span().execution_tags
         == [{"key": key, "value": value}] * MAX_TAGS  # noqa
     )
 
@@ -167,4 +165,4 @@ def __str__(self):
 
     assert add_execution_tag("key", ExceptionOnStr()) is False
     assert "Unable to add tag" in capsys.readouterr().out
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == []
+    assert SpansContainer.get_span().execution_tags == []

src/test/unit/wrappers/http/test_sync_http_wrappers.py

@@ -18,7 +18,6 @@
 from lumigo_tracer import lumigo_utils
 from lumigo_tracer.auto_tag import auto_tag_event
 from lumigo_tracer.lumigo_utils import (
-    EXECUTION_TAGS_KEY,
     DEFAULT_MAX_ENTRY_SIZE,
     Configuration,
     TRUNCATE_SUFFIX,
@@ -355,7 +354,7 @@ def lambda_test_function(event, context):
     result = lambda_test_function(api_gw_event(), context)
 
     assert result == "ret_value"
-    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == [
+    assert SpansContainer.get_span().execution_tags == [
         {"key": "Accept", "value": "application/json, text/plain, */*"}
     ]