add cycle detection

Author: liustve

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py

@@ -4,7 +4,7 @@
 import logging
 import os
 import re
-from logging import NOTSET, Logger, getLogger
+from logging import Logger, getLogger
 from typing import ClassVar, Dict, List, NamedTuple, Optional, Type, Union
 
 from importlib_metadata import version

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/exporter/otlp/aws/logs/aws_batch_log_record_processor.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
-from typing import List, Mapping, Optional, Sequence, cast
+from typing import Mapping, Optional, Sequence, cast
 
 from amazon.opentelemetry.distro.exporter.otlp.aws.logs.otlp_aws_logs_exporter import OTLPAwsLogExporter
 from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, attach, detach, set_value
@@ -98,11 +98,12 @@ def _export(self, batch_strategy: BatchLogExportStrategy) -> None:
                     _logger.exception("Exception while exporting logs: %s", exception)
                 detach(token)
 
-    def _estimate_log_size(self, log: LogData, depth: int = 3) -> int:
+    def _estimate_log_size(self, log: LogData, depth: int = 3) -> int:  # pylint: disable=too-many-branches
         """
         Estimates the size in bytes of a log by calculating the size of its body and its attributes
         and adding a buffer amount to account for other log metadata information.
         Will process complex log structures up to the specified depth limit.
+        Includes cycle detection to prevent processing the log content more than once.
         If the depth limit of the log structure is exceeded, returns the truncated calculation
         to everything up to that point.
 
@@ -114,14 +115,21 @@ def _estimate_log_size(self, log: LogData, depth: int = 3) -> int:
             int: The estimated size of the log object in bytes
         """
 
-        # Use a queue to prevent excessive recursive calls.
-        # We calculate based on the size of the log record body and attributes for the log.
-        queue: List[tuple[AnyValue, int]] = [(log.log_record.body, 0), (log.log_record.attributes, -1)]
+        # Queue contains tuples of (log_content, depth) where:
+        # - log_content is the current piece of log data being processed
+        # - depth tracks how many levels deep we've traversed to reach this content
+        # - body starts at depth 0 since it's an AnyValue object
+        # - Attributes start at depth -1 since it's a Mapping[str, AnyValue] - when traversed, we will
+        #   start processing its keys at depth 0
+        queue = [(log.log_record.body, 0), (log.log_record.attributes, -1)]
+
+        # Track visited complex log contents to avoid calculating the same one more than once
+        visited = set()
 
         size: int = self._BASE_LOG_BUFFER_BYTE_SIZE
 
         while queue:
-            new_queue: List[tuple[AnyValue, int]] = []
+            new_queue = []
 
             for data in queue:
                 # small optimization, can stop calculating the size once it reaches the 1 MB limit.
@@ -130,19 +138,30 @@ def _estimate_log_size(self, log: LogData, depth: int = 3) -> int:
 
                 next_val, current_depth = data
 
-                if isinstance(next_val, (str, bytes)):
-                    size += len(next_val)
+                if next_val is None:
                     continue
 
                 if isinstance(next_val, bool):
                     size += 4 if next_val else 5
                     continue
 
+                if isinstance(next_val, (str, bytes)):
+                    size += len(next_val)
+                    continue
+
                 if isinstance(next_val, (float, int)):
                     size += len(str(next_val))
                     continue
 
+                # next_val must be Sequence["AnyValue"] or Mapping[str, "AnyValue"],
                 if current_depth <= depth:
+                    obj_id = id(
+                        next_val
+                    )  # Guaranteed to be unique, see: https://www.w3schools.com/python/ref_func_id.asp
+                    if obj_id in visited:
+                        continue
+                    visited.add(obj_id)
+
                     if isinstance(next_val, Sequence):
                         for content in next_val:
                             new_queue.append((cast(AnyValue, content), current_depth + 1))

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/exporter/otlp/aws/logs/test_aws_batch_log_record_processor.py

@@ -24,96 +24,117 @@ def setUp(self):
         self.mock_exporter.export.return_value = LogExportResult.SUCCESS
 
         self.processor = AwsCloudWatchOtlpBatchLogRecordProcessor(exporter=self.mock_exporter)
-        self.max_log_size = self.processor._MAX_LOG_REQUEST_BYTE_SIZE
-        self.base_log_size = self.processor._BASE_LOG_BUFFER_BYTE_SIZE
 
     def test_process_log_data_nested_structure(self):
         """Tests that the processor correctly handles nested structures (dict/list)"""
-        message_size = 400
-        message = "X" * message_size
+        log_body = "X" * 400
+        log_key = "test"
+        log_depth = 2
 
-        nest_dict_log = self.generate_test_log_data(
-            log_body=message, attr_key="t", attr_val=message, log_body_depth=2, attr_depth=2, count=1, create_map=True
+        nested_dict_log = self.generate_test_log_data(
+            log_body=log_body, log_key=log_key, log_body_depth=log_depth, count=1, create_map=True
         )
-        nest_array_log = self.generate_test_log_data(
-            log_body=message, attr_key="t", attr_val=message, log_body_depth=2, attr_depth=2, count=1, create_map=False
+        nested_array_log = self.generate_test_log_data(
+            log_body=log_body, log_key=log_key, log_body_depth=log_depth, count=1, create_map=False
         )
 
-        expected_size = self.base_log_size + message_size * 2
+        expected_dict_size = len(log_key) * log_depth + len(log_body)
+        expected_array_size = len(log_body)
+
+        dict_size = self.processor._estimate_log_size(log=nested_dict_log[0], depth=log_depth)
+        array_size = self.processor._estimate_log_size(log=nested_array_log[0], depth=log_depth)
+
+        self.assertEqual(dict_size - self.processor._BASE_LOG_BUFFER_BYTE_SIZE, expected_dict_size)
+        self.assertEqual(array_size - self.processor._BASE_LOG_BUFFER_BYTE_SIZE, expected_array_size)
+
+    def test_process_log_data_with_attributes(self):
+        """Tests that the processor correctly handles both body and attributes"""
+        log_body = "test_body"
+        attr_key = "attr_key"
+        attr_value = "attr_value"
+
+        record = LogRecord(
+            timestamp=int(time.time_ns()),
+            trace_id=0x123456789ABCDEF0123456789ABCDEF0,
+            span_id=0x123456789ABCDEF0,
+            trace_flags=TraceFlags(1),
+            severity_text="INFO",
+            severity_number=SeverityNumber.INFO,
+            body=log_body,
+            attributes={attr_key: attr_value},
+        )
+        log_data = LogData(log_record=record, instrumentation_scope=InstrumentationScope("test-scope", "1.0.0"))
 
-        dict_size = self.processor._estimate_log_size(log=nest_dict_log[0], depth=2)
-        array_size = self.processor._estimate_log_size(log=nest_array_log[0], depth=2)
+        expected_size = len(log_body) + len(attr_key) + len(attr_value)
+        actual_size = self.processor._estimate_log_size(log_data)
 
-        # Asserting almost equal to account for dictionary keys in the Log object
-        self.assertAlmostEqual(dict_size, expected_size, delta=10)
-        self.assertAlmostEqual(array_size, expected_size, delta=10)
+        self.assertEqual(actual_size - self.processor._BASE_LOG_BUFFER_BYTE_SIZE, expected_size)
 
     def test_process_log_data_nested_structure_exceeds_depth(self):
         """Tests that the processor cuts off calculation for nested structure that exceeds the depth limit"""
-        calculated = "X" * 400
-        message = {"calculated": calculated, "truncated": {"truncated": {"test": "X" * self.max_log_size}}}
-
-        # *2 since we set this message in both body and attributes
-        expected_size = self.base_log_size + (len("calculated") + len(calculated) + len("truncated")) * 2
+        max_depth = 0
+        calculated_body = "X" * 400
+        log_body = {
+            "calculated": "X" * 400,
+            "restOfThisLogWillBeTruncated": {"truncated": {"test": "X" * self.processor._MAX_LOG_REQUEST_BYTE_SIZE}},
+        }
 
-        nest_dict_log = self.generate_test_log_data(
-            log_body=message, attr_key="t", attr_val=message, log_body_depth=3, attr_depth=3, count=1, create_map=True
-        )
-        nest_array_log = self.generate_test_log_data(
-            log_body=message, attr_key="t", attr_val=message, log_body_depth=3, attr_depth=3, count=1, create_map=False
+        expected_size = self.processor._BASE_LOG_BUFFER_BYTE_SIZE + (
+            len("calculated") + len(calculated_body) + len("restOfThisLogWillBeTruncated")
         )
 
-        # Only calculates log size of up to depth of 4
-        dict_size = self.processor._estimate_log_size(log=nest_dict_log[0], depth=4)
-        array_size = self.processor._estimate_log_size(log=nest_array_log[0], depth=4)
+        test_logs = self.generate_test_log_data(log_body=log_body, count=1)
+
+        # Only calculates log size of up to depth of 0
+        dict_size = self.processor._estimate_log_size(log=test_logs[0], depth=max_depth)
 
-        # Asserting almost equal to account for dictionary keys in the Log object body
-        self.assertAlmostEqual(dict_size, expected_size, delta=10)
-        self.assertAlmostEqual(array_size, expected_size, delta=10)
+        self.assertEqual(dict_size, expected_size)
 
     def test_process_log_data_nested_structure_size_exceeds_max_log_size(self):
         """Tests that the processor returns prematurely if the size already exceeds _MAX_LOG_REQUEST_BYTE_SIZE"""
-        # Should stop calculation at bigKey
-        message = {
-            "bigKey": "X" * (self.max_log_size),
-            "smallKey": "X" * (self.max_log_size * 10),
+        # Should stop calculation at bigKey + biggerKey and not calculate the content of biggerKey
+        log_body = {
+            "bigKey": "X" * (self.processor._MAX_LOG_REQUEST_BYTE_SIZE),
+            "biggerKey": "X" * (self.processor._MAX_LOG_REQUEST_BYTE_SIZE * 100),
         }
 
-        expected_size = self.base_log_size + self.max_log_size + len("bigKey")
-
-        nest_dict_log = self.generate_test_log_data(
-            log_body=message, attr_key="", attr_val="", log_body_depth=-1, attr_depth=-1, count=1, create_map=True
-        )
-        nest_array_log = self.generate_test_log_data(
-            log_body=message, attr_key="", attr_val="", log_body_depth=-1, attr_depth=-1, count=1, create_map=False
+        expected_size = (
+            self.processor._BASE_LOG_BUFFER_BYTE_SIZE
+            + self.processor._MAX_LOG_REQUEST_BYTE_SIZE
+            + len("bigKey")
+            + len("biggerKey")
         )
 
+        nest_dict_log = self.generate_test_log_data(log_body=log_body, count=1, create_map=True)
+        nest_array_log = self.generate_test_log_data(log_body=log_body, count=1, create_map=False)
+
         dict_size = self.processor._estimate_log_size(log=nest_dict_log[0])
         array_size = self.processor._estimate_log_size(log=nest_array_log[0])
 
-        self.assertAlmostEqual(dict_size, expected_size, delta=10)
-        self.assertAlmostEqual(array_size, expected_size, delta=10)
+        self.assertEqual(dict_size, expected_size)
+        self.assertEqual(array_size, expected_size)
 
     def test_process_log_data_primitive(self):
 
         primitives: List[AnyValue] = ["test", b"test", 1, 1.2, True, False, None]
         expected_sizes = [4, 4, 1, 3, 4, 5, 0]
 
         for index, primitive in enumerate(primitives):
-            log = self.generate_test_log_data(
-                log_body=primitive,
-                attr_key="",
-                attr_val="",
-                log_body_depth=-1,
-                attr_depth=-1,
-                count=1,
-            )
-
-            expected_size = self.base_log_size + expected_sizes[index]
+            log = self.generate_test_log_data(log_body=primitive, count=1)
+            expected_size = self.processor._BASE_LOG_BUFFER_BYTE_SIZE + expected_sizes[index]
             actual_size = self.processor._estimate_log_size(log[0])
-
             self.assertEqual(actual_size, expected_size)
 
+    def test_process_log_data_with_cycle(self):
+        """Test that processor handles processing logs with circular references only once"""
+        cyclic_dict: dict = {"data": "test"}
+        cyclic_dict["self_ref"] = cyclic_dict
+
+        log = self.generate_test_log_data(log_body=cyclic_dict, count=1)
+        expected_size = self.processor._BASE_LOG_BUFFER_BYTE_SIZE + len("data") + len("self_ref") + len("test")
+        actual_size = self.processor._estimate_log_size(log[0])
+        self.assertEqual(actual_size, expected_size)
+
     @patch(
         "amazon.opentelemetry.distro.exporter.otlp.aws.logs.aws_batch_log_record_processor.attach",
         return_value=MagicMock(),
@@ -124,9 +145,7 @@ def test_export_single_batch_under_size_limit(self, _, __, ___):
         """Tests that export is only called once if a single batch is under the size limit"""
         log_count = 10
         log_body = "test"
-        test_logs = self.generate_test_log_data(
-            log_body=log_body, attr_key="", attr_val="", log_body_depth=-1, attr_depth=-1, count=log_count
-        )
+        test_logs = self.generate_test_log_data(log_body=log_body, count=log_count)
         total_data_size = 0
 
         for log in test_logs:
@@ -153,9 +172,7 @@ def test_export_single_batch_all_logs_over_size_limit(self, _, __, ___):
         """Should make multiple export calls of batch size 1 to export logs of size > 1 MB."""
 
         large_log_body = "X" * (self.processor._MAX_LOG_REQUEST_BYTE_SIZE + 1)
-        test_logs = self.generate_test_log_data(
-            log_body=large_log_body, attr_key="", attr_val="", log_body_depth=-1, attr_depth=-1, count=15
-        )
+        test_logs = self.generate_test_log_data(log_body=large_log_body, count=15)
 
         for log in test_logs:
             self.processor._queue.appendleft(log)
@@ -178,26 +195,13 @@ def test_export_single_batch_all_logs_over_size_limit(self, _, __, ___):
     @patch("amazon.opentelemetry.distro.exporter.otlp.aws.logs.aws_batch_log_record_processor.set_value")
     def test_export_single_batch_some_logs_over_size_limit(self, _, __, ___):
         """Should make calls to export smaller sub-batch logs"""
-        large_log_body = "X" * (self.max_log_size + 1)
-        small_log_body = "X" * (self.max_log_size // 10 - self.base_log_size)
-
-        large_logs = self.generate_test_log_data(
-            log_body=large_log_body,
-            attr_key="",
-            attr_val="",
-            log_body_depth=-1,
-            attr_depth=-1,
-            count=3,
+        large_log_body = "X" * (self.processor._MAX_LOG_REQUEST_BYTE_SIZE + 1)
+        small_log_body = "X" * (
+            self.processor._MAX_LOG_REQUEST_BYTE_SIZE // 10 - self.processor._BASE_LOG_BUFFER_BYTE_SIZE
         )
 
-        small_logs = self.generate_test_log_data(
-            log_body=small_log_body,
-            attr_key="",
-            attr_val="",
-            log_body_depth=-1,
-            attr_depth=-1,
-            count=12,
-        )
+        large_logs = self.generate_test_log_data(log_body=large_log_body, count=3)
+        small_logs = self.generate_test_log_data(log_body=small_log_body, count=12)
 
         # 1st, 2nd, 3rd batch = size 1
         # 4th batch = size 10
@@ -230,36 +234,32 @@ def test_export_single_batch_some_logs_over_size_limit(self, _, __, ___):
     @staticmethod
     def generate_test_log_data(
         log_body,
-        attr_key,
-        attr_val,
-        log_body_depth=3,
-        attr_depth=3,
+        log_key="key",
+        log_body_depth=0,
         count=5,
         create_map=True,
-        instrumentation_scope=InstrumentationScope("test-scope", "1.0.0"),
     ) -> List[LogData]:
 
         def generate_nested_value(depth, value, create_map=True) -> AnyValue:
-            if depth < 0:
+            if depth <= 0:
                 return value
 
             if create_map:
-                return {"t": generate_nested_value(depth - 1, value, True)}
+                return {log_key: generate_nested_value(depth - 1, value, True)}
 
             return [generate_nested_value(depth - 1, value, False)]
 
         logs = []
 
-        for index in range(count):
+        for _ in range(count):
             record = LogRecord(
                 timestamp=int(time.time_ns()),
-                trace_id=int(f"0x{index + 1:032x}", 16),
-                span_id=int(f"0x{index + 1:016x}", 16),
+                trace_id=0x123456789ABCDEF0123456789ABCDEF0,
+                span_id=0x123456789ABCDEF0,
                 trace_flags=TraceFlags(1),
                 severity_text="INFO",
                 severity_number=SeverityNumber.INFO,
                 body=generate_nested_value(log_body_depth, log_body, create_map),
-                attributes={attr_key: generate_nested_value(attr_depth, attr_val, create_map)},
             )
 
             log_data = LogData(log_record=record, instrumentation_scope=InstrumentationScope("test-scope", "1.0.0"))