retry in case of error with edge (#149)

Author: saartochner-lumigo

src/lumigo_tracer/utils.py

@@ -215,23 +215,26 @@ def prepare_host(host):
     return host
 
 
-def report_json(region: Union[None, str], msgs: List[dict]) -> int:
+def report_json(region: Union[None, str], msgs: List[dict], should_retry: bool = True) -> int:
     """
     This function sends the information back to the edge.
 
     :param region: The region to use as default if not configured otherwise.
     :param msgs: the message to send.
+    :param should_retry: False to disable the default retry on unsuccessful sending
     :return: The duration of reporting (in milliseconds),
                 or 0 if we didn't send (due to configuration or fail).
     """
     global edge_connection
     get_logger().info(f"reporting the messages: {msgs[:10]}")
-    host = prepare_host(Configuration.host or EDGE_HOST.format(region=region))
-    duration = 0
-    if not edge_connection or edge_connection.host != host:
-        edge_connection = establish_connection(host)
-        if not edge_connection:
-            return duration
+    host = None
+    with lumigo_safe_execute("report json: establish connection"):
+        host = prepare_host(Configuration.host or EDGE_HOST.format(region=region))
+        duration = 0
+        if not edge_connection or edge_connection.host != host:
+            edge_connection = establish_connection(host)
+            if not edge_connection:
+                return duration
     if Configuration.should_report:
         try:
             prune_trace: bool = not os.environ.get("LUMIGO_PRUNE_TRACE_OFF", "").lower() == "true"
@@ -245,10 +248,12 @@ def report_json(region: Union[None, str], msgs: List[dict]) -> int:
             duration = int((time.time() - start_time) * 1000)
             get_logger().info(f"successful reporting, code: {getattr(response, 'code', 'unknown')}")
         except Exception as e:
-            get_logger().exception(
-                f"Could not report json to {host}. Retrying to establish connection.", exc_info=e
-            )
-            edge_connection = establish_connection(host)
+            if should_retry:
+                get_logger().exception(f"Could not report to {host}. Retrying.", exc_info=e)
+                edge_connection = establish_connection(host)
+                report_json(region, msgs, should_retry=False)
+            else:
+                get_logger().exception("Could not report: A span was lost.", exc_info=e)
     return duration
 
 

src/test/unit/test_main_utils.py

@@ -2,6 +2,8 @@
 from collections import OrderedDict
 from decimal import Decimal
 import datetime
+import http.client
+from mock import Mock
 
 import pytest
 from lumigo_tracer.utils import (
@@ -26,6 +28,7 @@
     lumigo_dumps,
     prepare_host,
     EDGE_PATH,
+    report_json,
 )
 import json
 
@@ -362,3 +365,18 @@ def test_get_timeout_buffer(remaining_time, conf, expected):
 )
 def test_prepare_host(arg, host):
     assert prepare_host(arg) == host
+
+
+@pytest.mark.parametrize(
+    "errors, final_log", [(ValueError, "ERROR"), ([ValueError, Mock()], "INFO")]
+)
+def test_report_json_retry(monkeypatch, reporter_mock, caplog, errors, final_log):
+    reporter_mock.side_effect = report_json
+    monkeypatch.setattr(Configuration, "host", "force_reconnect")
+    monkeypatch.setattr(Configuration, "should_report", True)
+    monkeypatch.setattr(http.client, "HTTPSConnection", Mock())
+    http.client.HTTPSConnection("force_reconnect").getresponse.side_effect = errors
+
+    report_json(None, [{"a": "b"}])
+
+    assert caplog.records[-1].levelname == final_log