chore(ci_visibility): split writer payloads (#15576)

## Description

Split payloads to the event intakes when they would exceed 4.5MB.

## Testing

Unit tests.

## Risks

None.

## Additional Notes

None.

Author: vitor-de-araujo

ddtrace/testing/internal/telemetry.py

@@ -177,7 +177,6 @@ def record_event_payload(
         payload_size: int,
         request_seconds: float,
         events_count: int,
-        serialization_seconds: float,
         error: t.Optional[ErrorType],
     ) -> None:
         tags = {"endpoint": endpoint}
@@ -186,11 +185,14 @@ def record_event_payload(
         self.add_count_metric("endpoint_payload.requests", 1, tags)
         self.add_distribution_metric("endpoint_payload.requests_ms", request_seconds * 1000, tags)
         self.add_distribution_metric("endpoint_payload.events_count", events_count, tags)
-        self.add_distribution_metric("endpoint_payload.events_serialization_ms", serialization_seconds * 1000, tags)
 
         if error:
             self.record_event_payload_error(endpoint, error)
 
+    def record_event_payload_serialization_seconds(self, endpoint: str, serialization_seconds: float) -> None:
+        tags = {"endpoint": endpoint}
+        self.add_distribution_metric("endpoint_payload.events_serialization_ms", serialization_seconds * 1000, tags)
+
     def record_event_payload_error(self, endpoint: str, error: ErrorType) -> None:
         # `endpoint_payload.requests_errors` accepts a different set of error types, so we need to convert them here.
         if error == ErrorType.TIMEOUT:

ddtrace/testing/internal/writer.py

@@ -35,9 +35,10 @@ def __init__(self) -> None:
         self.should_finish = threading.Event()
         self.flush_interval_seconds = 60
         self.events: t.List[Event] = []
+        # 4.5MB max uncompressed payload size, following <https://github.com/DataDog/datadog-ci-rb/pull/272>.
+        self.max_payload_size = int(4.5 * 1024 * 1024)
 
     def put_event(self, event: Event) -> None:
-        # TODO: compute/estimate payload size as events are inserted, and force a push once we reach a certain size.
         with self.lock:
             self.events.append(event)
 
@@ -80,6 +81,22 @@ def flush(self) -> None:
     def _send_events(self, events: t.List[Event]) -> None:
         pass
 
+    @abstractmethod
+    def _encode_events(self, events: t.List[Event]) -> bytes:
+        pass
+
+    def _split_pack_events(self, events: t.List[Event]) -> t.List[bytes]:
+        pack = self._encode_events(events)
+
+        if len(pack) > self.max_payload_size and len(events) > 1:
+            del pack
+            midpoint = len(events) // 2
+            packs = self._split_pack_events(events[0:midpoint])
+            packs += self._split_pack_events(events[midpoint:])
+            return packs
+        else:
+            return [pack]
+
 
 class TestOptWriter(BaseWriter):
     __test__ = False
@@ -122,27 +139,36 @@ def put_item(self, item: TestItem[t.Any, t.Any]) -> None:
         event = self.serializers[type(item)](item)
         self.put_event(event)
 
-    def _send_events(self, events: t.List[Event]) -> None:
+    def _encode_events(self, events: t.List[Event]) -> bytes:
         payload = {
             "version": 1,
             "metadata": self.metadata,
             "events": events,
         }
+        return msgpack_packb(payload)
+
+    def _send_events(self, events: t.List[Event]) -> None:
         with StopWatch() as serialization_time:
-            pack = msgpack_packb(payload)
+            packs = self._split_pack_events(events)
 
-        result = self.connector.request(
-            "POST", "/api/v2/citestcycle", data=pack, headers={"Content-Type": "application/msgpack"}, send_gzip=True
-        )
+        TelemetryAPI.get().record_event_payload_serialization_seconds("test_cycle", serialization_time.elapsed())
 
-        TelemetryAPI.get().record_event_payload(
-            endpoint="test_cycle",
-            payload_size=len(pack),
-            request_seconds=result.elapsed_seconds,
-            events_count=len(events),
-            serialization_seconds=serialization_time.elapsed(),
-            error=result.error_type,
-        )
+        for pack in packs:
+            result = self.connector.request(
+                "POST",
+                "/api/v2/citestcycle",
+                data=pack,
+                headers={"Content-Type": "application/msgpack"},
+                send_gzip=True,
+            )
+
+            TelemetryAPI.get().record_event_payload(
+                endpoint="test_cycle",
+                payload_size=len(pack),
+                request_seconds=result.elapsed_seconds,
+                events_count=len(events),
+                error=result.error_type,
+            )
 
 
 class TestCoverageWriter(BaseWriter):
@@ -169,35 +195,40 @@ def put_coverage(self, test_run: TestRun, coverage_bitmaps: t.Iterable[t.Tuple[s
         )
         self.put_event(event)
 
+    def _encode_events(self, events: t.List[Event]) -> bytes:
+        return msgpack_packb({"version": 2, "coverages": events})
+
     def _send_events(self, events: t.List[Event]) -> None:
         with StopWatch() as serialization_time:
-            pack = msgpack_packb({"version": 2, "coverages": events})
-
-        files = [
-            FileAttachment(
-                name="coverage1",
-                filename="coverage1.msgpack",
-                content_type="application/msgpack",
-                data=pack,
-            ),
-            FileAttachment(
-                name="event",
-                filename="event.json",
-                content_type="application/json",
-                data=b'{"dummy":true}',
-            ),
-        ]
-
-        result = self.connector.post_files("/api/v2/citestcov", files=files, send_gzip=True)
-
-        TelemetryAPI.get().record_event_payload(
-            endpoint="code_coverage",
-            payload_size=len(pack),
-            request_seconds=result.elapsed_seconds,
-            events_count=len(events),
-            serialization_seconds=serialization_time.elapsed(),
-            error=result.error_type,
-        )
+            packs = self._split_pack_events(events)
+
+        TelemetryAPI.get().record_event_payload_serialization_seconds("code_coverage", serialization_time.elapsed())
+
+        for pack in packs:
+            files = [
+                FileAttachment(
+                    name="coverage1",
+                    filename="coverage1.msgpack",
+                    content_type="application/msgpack",
+                    data=pack,
+                ),
+                FileAttachment(
+                    name="event",
+                    filename="event.json",
+                    content_type="application/json",
+                    data=b'{"dummy":true}',
+                ),
+            ]
+
+            result = self.connector.post_files("/api/v2/citestcov", files=files, send_gzip=True)
+
+            TelemetryAPI.get().record_event_payload(
+                endpoint="code_coverage",
+                payload_size=len(pack),
+                request_seconds=result.elapsed_seconds,
+                events_count=len(events),
+                error=result.error_type,
+            )
 
 
 def serialize_test_run(test_run: TestRun) -> Event:

tests/testing/internal/test_telemetry.py

@@ -274,7 +274,6 @@ def test_record_event_payload_ok(self, telemetry_api: TelemetryAPI) -> None:
             payload_size=613,
             request_seconds=3.14,
             events_count=42,
-            serialization_seconds=0.5,
             error=None,
         )
 
@@ -285,7 +284,6 @@ def test_record_event_payload_ok(self, telemetry_api: TelemetryAPI) -> None:
             call(CIVISIBILITY, "endpoint_payload.bytes", 613, (("endpoint", "test_cycle"),)),
             call(CIVISIBILITY, "endpoint_payload.requests_ms", 3140, (("endpoint", "test_cycle"),)),
             call(CIVISIBILITY, "endpoint_payload.events_count", 42, (("endpoint", "test_cycle"),)),
-            call(CIVISIBILITY, "endpoint_payload.events_serialization_ms", 500, (("endpoint", "test_cycle"),)),
         ]
 
     @pytest.mark.parametrize(
@@ -307,7 +305,6 @@ def test_record_event_payload_error(
             payload_size=613,
             request_seconds=3.14,
             events_count=42,
-            serialization_seconds=0.5,
             error=http_error_type,
         )
 
@@ -327,6 +324,12 @@ def test_record_event_payload_error(
             call(CIVISIBILITY, "endpoint_payload.bytes", 613, (("endpoint", "test_cycle"),)),
             call(CIVISIBILITY, "endpoint_payload.requests_ms", 3140, (("endpoint", "test_cycle"),)),
             call(CIVISIBILITY, "endpoint_payload.events_count", 42, (("endpoint", "test_cycle"),)),
+        ]
+
+    def test_record_event_payload_serialization_seconds(self, telemetry_api: TelemetryAPI) -> None:
+        telemetry_api.record_event_payload_serialization_seconds("test_cycle", 0.5)
+
+        assert telemetry_api.writer.add_distribution_metric.call_args_list == [
             call(CIVISIBILITY, "endpoint_payload.events_serialization_ms", 500, (("endpoint", "test_cycle"),)),
         ]
 

tests/testing/internal/test_writer.py

@@ -1,6 +1,7 @@
 """Tests for ddtrace.testing.internal.writer module."""
 
 from unittest.mock import Mock
+from unittest.mock import call
 from unittest.mock import patch
 
 from ddtrace.testing.internal.http import BackendConnectorAgentlessSetup
@@ -10,6 +11,7 @@
 from ddtrace.testing.internal.test_data import TestSession
 from ddtrace.testing.internal.test_data import TestStatus
 from ddtrace.testing.internal.test_data import TestSuite
+from ddtrace.testing.internal.tracer_api import msgpack_packb
 from ddtrace.testing.internal.writer import Event
 from ddtrace.testing.internal.writer import TestCoverageWriter
 from ddtrace.testing.internal.writer import TestOptWriter
@@ -101,6 +103,48 @@ def test_send_events(self, mock_packb: Mock, mock_backend_connector: Mock) -> No
             send_gzip=True,
         )
 
+    @patch("ddtrace.testing.internal.http.BackendConnector")
+    def test_split_events(self, mock_backend_connector: Mock) -> None:
+        """Test sending events to backend."""
+        mock_connector = Mock()
+        mock_backend_connector.return_value = mock_connector
+        mock_connector.request.return_value = BackendResult(
+            response=Mock(status=200), response_length=42, elapsed_seconds=1.2
+        )
+
+        writer = TestOptWriter(BackendConnectorAgentlessSetup(site="test", api_key="key"))
+        writer.metadata = {"*": {"foo": "bar"}}
+        writer.max_payload_size = 70
+
+        events = [Event(type="test1"), Event(type="test2"), Event(type="test3")]
+
+        writer._send_events(events)
+
+        expected_pack_1 = msgpack_packb(
+            {"version": 1, "metadata": {"*": {"foo": "bar"}}, "events": [{"type": "test1"}]}
+        )
+        expected_pack_2 = msgpack_packb(
+            {"version": 1, "metadata": {"*": {"foo": "bar"}}, "events": [{"type": "test2"}, {"type": "test3"}]}
+        )
+
+        # Check HTTP request
+        assert mock_connector.request.call_args_list == [
+            call(
+                "POST",
+                "/api/v2/citestcycle",
+                data=expected_pack_1,
+                headers={"Content-Type": "application/msgpack"},
+                send_gzip=True,
+            ),
+            call(
+                "POST",
+                "/api/v2/citestcycle",
+                data=expected_pack_2,
+                headers={"Content-Type": "application/msgpack"},
+                send_gzip=True,
+            ),
+        ]
+
 
 class TestTestCoverageWriter:
     """Tests for TestCoverageWriter class."""