Ensure thread safety + parsing errors don't crash instrumentation

Author: majanjua-amzn

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

@@ -1,6 +1,9 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 # Modifications Copyright The OpenTelemetry Authors. Licensed under the Apache License 2.0 License.
+
+# pylint: disable=too-many-lines
+
 import logging
 import os
 import re
@@ -413,7 +416,8 @@ def _customize_sampler(sampler: Sampler) -> Sampler:
 
         try:
             parsed_config = _parse_config_string(config)
-        except ValueError as error:
+        # pylint: disable=broad-exception-caught
+        except Exception as error:
             _logger.warning("Failed to parse adaptive sampling configuration: %s", str(error))
 
         if parsed_config is not None:
@@ -924,7 +928,7 @@ def _create_aws_otlp_exporter(endpoint: str, service: str, region: str):
         _logger.error("Failed to create AWS OTLP exporter: %s", errors)
         return None
 
-
+# pylint: disable=too-many-return-statements,too-many-branches
 def _parse_config_string(config: str) -> Optional[_AWSXRayAdaptiveSamplingConfig]:
     if config is None:
         return None
@@ -935,54 +939,71 @@ def _parse_config_string(config: str) -> Optional[_AWSXRayAdaptiveSamplingConfig
         try:
             config = path.read_text(encoding="utf-8")
         except IOError as err:
-            raise ValueError(f"Failed to read adaptive sampling configuration file: {err}") from err
+            _logger.warning("Failed to read adaptive sampling configuration file: %s", err)
+            return None
     elif config.endswith(".yml") or config.endswith(".yaml"):
-        raise ValueError("Adaptive sampling configuration file must be a YAML file")
+        _logger.warning("Adaptive sampling configuration file must be a YAML file")
+        return None
     else:
         _logger.debug("Adaptive sampling configuration is not a file path, assuming it's a YAML string")
 
     # Parse YAML config
-    config_map = yaml.safe_load(config)
+    config_map = None
+    try:
+        config_map = yaml.safe_load(config)
+    # pylint: disable=broad-exception-caught
+    except Exception as exception:
+        _logger.warning("Adaptive sampling configuration must be a valid YAML mapping: %s", exception)
     if not isinstance(config_map, dict):
-        raise ValueError("Adaptive sampling configuration must be a valid YAML mapping")
+        _logger.warning("Adaptive sampling configuration must be a valid YAML mapping")
+        return None
 
     # Ensure only relevant data is in the YAML configuration
     for key in config_map:
         if key not in ["version", "anomalyConditions", "anomalyCaptureLimit"]:
-            raise ValueError(f"Invalid key in adaptive sampling configuration: {key}")
+            _logger.warning("Invalid key in adaptive sampling configuration: %s", key)
+            return None
 
     version_obj = config_map.get("version")
     if version_obj is None:
-        raise ValueError("Missing required 'version' field in adaptive sampling configuration")
+        _logger.warning("Missing required 'version' field in adaptive sampling configuration")
+        return None
 
-    config_version = float(version_obj)
-    if config_version < 1.0 or config_version >= 2.0:
-        raise ValueError(
-            f"Incompatible adaptive sampling config version: {config_version}. "
-            "This version of the AWS X-Ray remote sampler only supports version 1.X."
-        )
+    try:
+        config_version = float(version_obj)
+        if config_version < 1.0 or config_version >= 2.0:
+            _logger.warning(
+                "Incompatible adaptive sampling config version: %s. "
+                "This version of the AWS X-Ray remote sampler only supports version 1.X.",
+                config_version,
+            )
+            return None
 
-    # Parse anomaly conditions
-    anomaly_conditions = None
-    if "anomalyConditions" in config_map:
-        anomaly_conditions = [
-            _AnomalyConditions(
-                error_code_regex=cond.get("errorCodeRegex"),
-                operations=cond.get("operations"),
-                high_latency_ms=cond.get("highLatencyMs"),
-                usage=_UsageType(cond["usage"]) if "usage" in cond else None,
+        # Parse anomaly conditions
+        anomaly_conditions = None
+        if "anomalyConditions" in config_map:
+            anomaly_conditions = [
+                _AnomalyConditions(
+                    error_code_regex=cond.get("errorCodeRegex"),
+                    operations=cond.get("operations"),
+                    high_latency_ms=cond.get("highLatencyMs"),
+                    usage=_UsageType(cond["usage"]) if "usage" in cond else None,
+                )
+                for cond in config_map["anomalyConditions"]
+            ]
+
+        # Parse anomaly capture limit
+        anomaly_capture_limit = None
+        if "anomalyCaptureLimit" in config_map:
+            anomaly_capture_limit_data = config_map["anomalyCaptureLimit"]
+            anomaly_capture_limit = _AnomalyCaptureLimit(
+                anomaly_traces_per_second=anomaly_capture_limit_data["anomalyTracesPerSecond"]
             )
-            for cond in config_map["anomalyConditions"]
-        ]
 
-    # Parse anomaly capture limit
-    anomaly_capture_limit = None
-    if "anomalyCaptureLimit" in config_map:
-        anomaly_capture_limit_data = config_map["anomalyCaptureLimit"]
-        anomaly_capture_limit = _AnomalyCaptureLimit(
-            anomaly_traces_per_second=anomaly_capture_limit_data["anomalyTracesPerSecond"]
+        return _AWSXRayAdaptiveSamplingConfig(
+            version=config_version, anomaly_conditions=anomaly_conditions, anomaly_capture_limit=anomaly_capture_limit
         )
+    except ValueError as err:
+        _logger.warning("Failed to load AWS X-Ray adaptive sampling configuration: %s", err)
 
-    return _AWSXRayAdaptiveSamplingConfig(
-        version=config_version, anomaly_conditions=anomaly_conditions, anomaly_capture_limit=anomaly_capture_limit
-    )
+    return None

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/sampler/_rule_cache.py

@@ -59,6 +59,7 @@ def __init__(
         self._adaptive_sampling_rule_exists = False
         self._adaptive_sampling_config: Optional[_AWSXRayAdaptiveSamplingConfig] = None
         self._anomaly_capture_rate_limiter: Optional[_RateLimiter] = None
+        self._trace_usage_cache_lock = Lock()
         self._trace_usage_cache: TTLCache[int, _UsageType] = TTLCache(
             maxsize=TRACE_USAGE_CACHE_MAX_SIZE,
             ttl=TRACE_USAGE_CACHE_TTL_SECONDS,
@@ -151,8 +152,12 @@ def adapt_sampling(self, span: ReadableSpan, span_batcher: Callable[[ReadableSpa
         should_capture_anomaly_span = result.should_capture_anomaly_span
 
         trace_id: int = span.context.trace_id
-        existing_usage: _UsageType = self._trace_usage_cache.get(trace_id)
-        is_new_trace: bool = existing_usage is None
+        is_new_trace: bool = False
+        with self._trace_usage_cache_lock:
+            existing_usage: _UsageType = self._trace_usage_cache.get(trace_id)
+            is_new_trace = existing_usage is None
+            if existing_usage is None:
+                self._trace_usage_cache[trace_id] = _UsageType.NEITHER
 
         # Anomaly Capture
         is_span_captured = False
@@ -275,25 +280,26 @@ def __is_anomaly(self, span: ReadableSpan) -> "_AnomalyDetectionResult":
     def __update_trace_usage_cache(
         self, trace_id: int, is_span_captured: bool, is_counted_as_anomaly_for_boost: bool
     ) -> None:
-        existing_usage = self._trace_usage_cache.get(trace_id)
+        with self._trace_usage_cache_lock:
+            existing_usage = self._trace_usage_cache.get(trace_id)
 
-        # Any interaction with a cache entry will reset the expiration timer of that entry
-        if is_span_captured and is_counted_as_anomaly_for_boost:
-            self._trace_usage_cache[trace_id] = _UsageType.BOTH
-        elif is_span_captured:
-            if _UsageType.is_used_for_boost(existing_usage):
-                self._trace_usage_cache[trace_id] = _UsageType.BOTH
-            else:
-                self._trace_usage_cache[trace_id] = _UsageType.ANOMALY_TRACE_CAPTURE
-        elif is_counted_as_anomaly_for_boost:
-            if _UsageType.is_used_for_anomaly_trace_capture(existing_usage):
+            # Any interaction with a cache entry will reset the expiration timer of that entry
+            if is_span_captured and is_counted_as_anomaly_for_boost:
                 self._trace_usage_cache[trace_id] = _UsageType.BOTH
+            elif is_span_captured:
+                if _UsageType.is_used_for_boost(existing_usage):
+                    self._trace_usage_cache[trace_id] = _UsageType.BOTH
+                else:
+                    self._trace_usage_cache[trace_id] = _UsageType.ANOMALY_TRACE_CAPTURE
+            elif is_counted_as_anomaly_for_boost:
+                if _UsageType.is_used_for_anomaly_trace_capture(existing_usage):
+                    self._trace_usage_cache[trace_id] = _UsageType.BOTH
+                else:
+                    self._trace_usage_cache[trace_id] = _UsageType.SAMPLING_BOOST
+            elif existing_usage is not None:
+                self._trace_usage_cache[trace_id] = existing_usage
             else:
-                self._trace_usage_cache[trace_id] = _UsageType.SAMPLING_BOOST
-        elif existing_usage is not None:
-            self._trace_usage_cache[trace_id] = existing_usage
-        else:
-            self._trace_usage_cache[trace_id] = _UsageType.NEITHER
+                self._trace_usage_cache[trace_id] = _UsageType.NEITHER
 
     def update_sampling_rules(self, new_sampling_rules: List[_SamplingRule]) -> None:
         new_sampling_rules.sort()

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py

@@ -351,19 +351,46 @@ def test_parse_adaptive_sampling_config_null(self):
         self.assertIsNone(result)
 
     def test_parse_adaptive_sampling_config_missing_version(self):
-        """Tests that _parse_config_string raises exception for missing version"""
-        with self.assertRaises(Exception):
-            _parse_config_string("")
+        """Tests that _parse_config_string returns None for missing version and logs warning"""
+        with patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator._logger") as mock_logger:
+            result = _parse_config_string("anomalyConditions: []")
+            self.assertIsNone(result)
+            mock_logger.warning.assert_called_with(
+                "Missing required 'version' field in adaptive sampling configuration"
+            )
 
     def test_parse_adaptive_sampling_config_unsupported_version(self):
-        """Tests that _parse_config_string raises exception for unsupported version"""
-        with self.assertRaises(ValueError):
-            _parse_config_string("{version: 5000.1}")
+        """Tests that _parse_config_string returns None for unsupported version and logs warning"""
+        with patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator._logger") as mock_logger:
+            result = _parse_config_string("{version: 5000.1}")
+            self.assertIsNone(result)
+            mock_logger.warning.assert_called_with(
+                "Incompatible adaptive sampling config version: %s. "
+                "This version of the AWS X-Ray remote sampler only supports version 1.X.",
+                5000.1,
+            )
 
     def test_parse_adaptive_sampling_config_invalid_yaml(self):
-        """Tests that _parse_config_string raises exception for invalid YAML"""
-        with self.assertRaises(Exception):
-            _parse_config_string("{version: 1, invalid: yaml: structure}")
+        """Tests that _parse_config_string returns None for invalid YAML and logs warning"""
+        with patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator._logger") as mock_logger_1:
+            result = _parse_config_string("{version: 1, invalid: yaml: structure}")
+            self.assertIsNone(result)
+            mock_logger_1.warning.assert_called()
+            warning_calls = [str(call) for call in mock_logger_1.warning.call_args_list]
+            self.assertTrue(
+                any("must be a valid YAML mapping" in call for call in warning_calls),
+                f"Expected warning about invalid YAML, got: {warning_calls}",
+            )
+
+        with patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator._logger") as mock_logger_2:
+            result = _parse_config_string("{version: 1, {anomalyCaptureLimit: {anomalyTracesPerSecond: 1}}}")
+            self.assertIsNone(result)
+            mock_logger_2.warning.assert_called()
+            warning_calls = [str(call) for call in mock_logger_2.warning.call_args_list]
+            self.assertTrue(
+                any("must be a valid YAML mapping" in call for call in warning_calls),
+                f"Expected warning about invalid YAML, got: {warning_calls}",
+            )
 
     def test_parse_adaptive_sampling_config_from_file_valid(self):
         """Tests that _parse_config_string correctly parses a valid YAML file"""
@@ -376,20 +403,29 @@ def test_parse_adaptive_sampling_config_from_file_valid(self):
         self.assertEqual(config.anomaly_capture_limit.anomaly_traces_per_second, 10)
 
     def test_parse_adaptive_sampling_config_from_file_invalid(self):
-        """Tests that _parse_config_string raises exception for invalid YAML file"""
+        """Tests that _parse_config_string returns None for invalid YAML file and logs warning"""
         import os
 
         config_path = os.path.join(DATA_DIR, "adaptive-sampling-config-invalid.yaml")
-        with self.assertRaises(Exception):
-            _parse_config_string(config_path)
+        with patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator._logger") as mock_logger:
+            result = _parse_config_string(config_path)
+            self.assertIsNone(result)
+            mock_logger.warning.assert_called()
+            warning_calls = [str(call) for call in mock_logger.warning.call_args_list]
+            self.assertTrue(
+                any("Failed to load AWS X-Ray adaptive sampling configuration" in call for call in warning_calls),
+                f"Expected warning about failed configuration loading, got: {warning_calls}",
+            )
 
     def test_parse_adaptive_sampling_config_from_file_non_existant(self):
-        """Tests that _parse_config_string raises exception for non-existent YAML file"""
+        """Tests that _parse_config_string returns None for non-existent YAML file and logs warning"""
         import os
 
         config_path = os.path.join(DATA_DIR, "done.yaml")
-        with self.assertRaises(ValueError):
-            _parse_config_string(config_path)
+        with patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator._logger") as mock_logger:
+            result = _parse_config_string(config_path)
+            self.assertIsNone(result)
+            mock_logger.warning.assert_called_with("Adaptive sampling configuration file must be a YAML file")
 
     def test_customize_span_exporter(self):
         mock_exporter: SpanExporter = MagicMock(spec=OTLPSpanExporter)