RD-5479 (#199)

* RD-5479
create extensions api in tracer

Author: OrrLumigo

src/lumigo_tracer/lumigo_utils.py

@@ -16,6 +16,7 @@
 from base64 import b64encode
 import inspect
 import traceback
+from pathlib import Path
 
 LUMIGO_DOMAINS_SCRUBBER_KEY = "LUMIGO_DOMAINS_SCRUBBER"
 
@@ -67,6 +68,8 @@
 NUMBER_OF_SPANS_IN_REPORT_OPTIMIZATION = 200
 DEFAULT_KEY_DEPTH = 4
 LUMIGO_TOKEN_KEY = "LUMIGO_TRACER_TOKEN"
+LUMIGO_USE_TRACER_EXTENSION = "LUMIGO_USE_TRACER_EXTENSION"
+EXTENSION_DIR = "/tmp/lumigo-spans"
 KILL_SWITCH = "LUMIGO_SWITCH_OFF"
 ERROR_SIZE_LIMIT_MULTIPLIER = 2
 CHINA_REGION = "cn-northwest-1"
@@ -81,6 +84,10 @@
 internal_error_already_logged = False
 
 
+def should_use_tracer_extension() -> bool:
+    return (os.environ.get(LUMIGO_USE_TRACER_EXTENSION) or "false").lower() == "true"
+
+
 def get_region() -> str:
     return os.environ.get("AWS_REGION") or "UNKNOWN"
 
@@ -305,6 +312,10 @@ def report_json(region: Optional[str], msgs: List[dict], should_retry: bool = Tr
     except Exception as e:
         get_logger().exception("Failed to create request: A span was lost.", exc_info=e)
         return 0
+    if should_use_tracer_extension():
+        with lumigo_safe_execute("report json file: writing spans to file"):
+            write_spans_to_file(to_send)
+        return 0
     if region == CHINA_REGION:
         return _publish_spans_to_kinesis(to_send, CHINA_REGION)
     host = None
@@ -337,6 +348,15 @@ def report_json(region: Optional[str], msgs: List[dict], should_retry: bool = Tr
     return duration
 
 
+def write_spans_to_file(to_send: bytes) -> None:
+    file_name = f"{hashlib.md5(to_send).hexdigest()}_single"
+    Path(EXTENSION_DIR).mkdir(parents=True, exist_ok=True)
+    file_path = os.path.join(EXTENSION_DIR, file_name)
+    get_logger().info(f"writing spans to file {file_path}")
+    with open(file_path, "wb") as spans_file:
+        spans_file.write(to_send)
+
+
 def _publish_spans_to_kinesis(to_send: bytes, region: str) -> int:
     start_time = time.time()
     try:

src/test/conftest.py

@@ -1,13 +1,15 @@
 import builtins
 import logging
+import os
+import shutil
 from types import SimpleNamespace
 
-from lumigo_tracer import lumigo_utils
-from lumigo_tracer.spans_container import SpansContainer
 import mock
 import pytest
 
+from lumigo_tracer import lumigo_utils
 from lumigo_tracer.lumigo_utils import Configuration, get_omitting_regex, get_logger, get_edge_host
+from lumigo_tracer.spans_container import SpansContainer
 from lumigo_tracer.wrappers.http.http_data_classes import HttpState
 
 
@@ -83,3 +85,10 @@ def capture_all_logs(caplog):
 @pytest.fixture
 def context():
     return SimpleNamespace(aws_request_id="1234", get_remaining_time_in_millis=lambda: 1000 * 2)
+
+
+@pytest.fixture(autouse=True)
+def extension_clean():
+    yield
+    if os.path.exists("/tmp/lumigo-spans"):
+        shutil.rmtree("/tmp/lumigo-spans")

src/test/unit/test_lumigo_utils.py

@@ -1,5 +1,6 @@
 import importlib.util
 import inspect
+import hashlib
 import logging
 from collections import OrderedDict
 from decimal import Decimal
@@ -442,6 +443,26 @@ def test_get_edge_host(arg, host, monkeypatch):
     assert get_edge_host("region") == host
 
 
+def test_report_json_extension(monkeypatch, reporter_mock):
+    monkeypatch.setattr(Configuration, "should_report", True)
+    monkeypatch.setenv("LUMIGO_USE_TRACER_EXTENSION", "TRUE")
+    single = [{"a": "b"}]
+    to_send = _create_request_body(single, True).encode()
+    md5str = str(hashlib.md5(to_send).hexdigest())
+    file_name = f"{md5str}_single"
+    file_path = f"/tmp/lumigo-spans/{file_name}"
+    asserting_extension(file_path, single)
+    # test that same file doesnt cause error
+    asserting_extension(file_path, single)
+
+
+def asserting_extension(file_path, single):
+    duration = report_json(None, [{"a": "b"}])
+    span_from_file = json.load(open(file_path, "r"))
+    assert duration == 0
+    assert span_from_file == single
+
+
 @pytest.mark.parametrize(
     "errors, final_log", [(ValueError, "ERROR"), ([ValueError, Mock()], "INFO")]
 )