RD-4202 - prevent collection of lumigo events

src/lumigo_tracer/lumigo_utils.py

@@ -246,15 +246,16 @@ def establish_connection(host):
     return None
 
 
-def prepare_host(host):
+def get_edge_host(region: Optional[str] = None) -> str:
+    host = Configuration.host or EDGE_HOST.format(region=region or get_region())
     if host.startswith(HTTPS_PREFIX):
         host = host[len(HTTPS_PREFIX) :]  # noqa: E203
     if host.endswith(EDGE_PATH):
         host = host[: -len(EDGE_PATH)]
     return host
 
 
-def report_json(region: Union[None, str], msgs: List[dict], should_retry: bool = True) -> int:
+def report_json(region: Optional[str], msgs: List[dict], should_retry: bool = True) -> int:
     """
     This function sends the information back to the edge.
 
@@ -278,7 +279,7 @@ def report_json(region: Union[None, str], msgs: List[dict], should_retry: bool =
     host = None
     global edge_connection
     with lumigo_safe_execute("report json: establish connection"):
-        host = prepare_host(Configuration.host or EDGE_HOST.format(region=region))
+        host = get_edge_host(region)
         duration = 0
         if not edge_connection or edge_connection.host != host:
             edge_connection = establish_connection(host)
@@ -626,3 +627,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 get_region() -> str:
+    return os.environ.get("AWS_REGION") or "UNKNOWN"

src/lumigo_tracer/spans_container.py

@@ -22,6 +22,7 @@
     _is_span_has_error,
     create_step_function_span,
     get_current_ms_time,
+    get_region,
 )
 from lumigo_tracer import lumigo_utils
 from lumigo_tracer.parsing_utils import parse_trace_id, safe_split_get, recursive_json_join
@@ -115,10 +116,10 @@ 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]
+        self.span_ids_to_send.clear()
         if Configuration.send_only_if_error:
             to_send.append(self._generate_start_span())
         lumigo_utils.report_json(region=self.region, msgs=to_send)
-        self.span_ids_to_send.clear()
 
     def start_timeout_timer(self, context=None) -> None:
         if Configuration.timeout_timer:
@@ -158,7 +159,9 @@ def update_event_end_time(self) -> None:
         This function assumes synchronous execution - we update the last http event.
         """
         if self.spans:
-            self.spans[-1]["ended"] = get_current_ms_time()
+            span = self.spans[-1]
+            span["ended"] = get_current_ms_time()
+            self.span_ids_to_send.add(span["id"])
 
     def update_event_times(
         self, start_time: Optional[datetime] = None, end_time: Optional[datetime] = None
@@ -290,7 +293,7 @@ def create_span(cls, event=None, context=None, force=False) -> "SpansContainer":
             started=get_current_ms_time(),
             name=os.environ.get("AWS_LAMBDA_FUNCTION_NAME"),
             runtime=os.environ.get("AWS_EXECUTION_ENV"),
-            region=os.environ.get("AWS_REGION"),
+            region=get_region(),
             memory_allocated=os.environ.get("AWS_LAMBDA_FUNCTION_MEMORY_SIZE"),
             log_stream_name=os.environ.get("AWS_LAMBDA_LOG_STREAM_NAME"),
             log_group_name=os.environ.get("AWS_LAMBDA_LOG_GROUP_NAME"),

src/lumigo_tracer/wrappers/http/sync_http_wrappers.py

@@ -14,6 +14,7 @@
     lumigo_dumps,
     get_size_upper_bound,
     is_error_code,
+    get_edge_host,
 )
 from lumigo_tracer.spans_container import SpansContainer
 from lumigo_tracer.wrappers.http.http_data_classes import HttpRequest, HttpState
@@ -29,10 +30,19 @@
 HookedData = namedtuple("HookedData", ["headers", "path"])
 
 
+def is_lumigo_edge(host: Optional[str]) -> bool:
+    if host and get_edge_host() in host:
+        get_logger().info("Dropping Lumigo event to edge")
+        return True
+    return False
+
+
 def add_request_event(parse_params: HttpRequest):
     """
     This function parses an request event and add it to the span.
     """
+    if is_lumigo_edge(parse_params.host):
+        return
     parser = get_parser(parse_params.host)()
     msg = parser.parse_request(parse_params)
     HttpState.previous_request = parse_params
@@ -47,6 +57,8 @@ def add_unparsed_request(parse_params: HttpRequest):
     In that case, we will consider it as a continuance of the previous request if they got the same url,
         and we didn't get any answer yet.
     """
+    if is_lumigo_edge(parse_params.host):
+        return
     last_event = SpansContainer.get_span().get_last_span()
     if last_event:
         if last_event and last_event.get("type") == HTTP_TYPE and HttpState.previous_request:
@@ -69,6 +81,8 @@ def update_event_response(
                         the aggregated response body
     This function assumes synchronous execution - we update the last http event.
     """
+    if is_lumigo_edge(host):
+        return
     last_event = SpansContainer.get_span().pop_last_span()
     if last_event:
         http_info = last_event.get("info", {}).get("httpInfo", {})

src/test/unit/test_lumigo_utils.py

@@ -27,7 +27,7 @@
     SKIP_SCRUBBING_KEYS,
     get_timeout_buffer,
     lumigo_dumps,
-    prepare_host,
+    get_edge_host,
     EDGE_PATH,
     report_json,
     is_kill_switch_on,
@@ -370,8 +370,9 @@ def test_get_timeout_buffer(remaining_time, conf, expected):
     ["arg", "host"],
     [("https://a.com", "a.com"), (f"https://b.com{EDGE_PATH}", "b.com"), ("h.com", "h.com")],
 )
-def test_prepare_host(arg, host):
-    assert prepare_host(arg) == host
+def test_get_edge_host(arg, host, monkeypatch):
+    monkeypatch.setattr(Configuration, "host", arg)
+    assert get_edge_host("region") == host
 
 
 @pytest.mark.parametrize(

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

@@ -19,7 +19,11 @@
 from lumigo_tracer.wrappers.http.http_parser import Parser
 from lumigo_tracer.spans_container import SpansContainer
 from lumigo_tracer.wrappers.http.http_data_classes import HttpState, HttpRequest
-from lumigo_tracer.wrappers.http.sync_http_wrappers import add_request_event, update_event_response
+from lumigo_tracer.wrappers.http.sync_http_wrappers import (
+    add_request_event,
+    update_event_response,
+    is_lumigo_edge,
+)
 
 
 def test_lambda_wrapper_http(context):
@@ -479,3 +483,9 @@ def lambda_test_function(event, context):
 
     assert len(request_with_error["body"]) > len(request_no_error["body"])
     assert request_with_error["body"] == json.dumps(d)
+
+
+@pytest.mark.parametrize("host, is_lumigo", [("https://lumigo.io", True), ("google.com", False)])
+def test_is_lumigo_edge(host, is_lumigo, monkeypatch):
+    monkeypatch.setattr(Configuration, "host", "lumigo.io")
+    assert is_lumigo_edge(host) == is_lumigo