customize Starlette Instrumentor middleware to exclude noisy receive and send spans

Author: liustve

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_instrumentation_patch.py

@@ -67,7 +67,7 @@ def apply_instrumentation_patches() -> None:
         from amazon.opentelemetry.distro.patches._starlette_patches import _apply_starlette_instrumentation_patches
 
         # Starlette auto-instrumentation v0.54b includes a strict dependency version check
-        # This restriction was removed in v1.34.0/0.55b0. Applying temporary patch for Genesis launch
+        # This restriction was removed in v1.34.0/0.55b0. Applying temporary patch for Bedrock AgentCore launch
         # TODO: Remove patch after syncing with upstream v1.34.0 or later
         _apply_starlette_instrumentation_patches()
 

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_starlette_patches.py

@@ -4,6 +4,8 @@
 from logging import Logger, getLogger
 from typing import Collection
 
+from amazon.opentelemetry.distro._utils import AGENT_OBSERVABILITY_ENABLED
+
 _logger: Logger = getLogger(__name__)
 
 
@@ -18,6 +20,7 @@ def _apply_starlette_instrumentation_patches() -> None:
     """
     try:
         # pylint: disable=import-outside-toplevel
+        from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware
         from opentelemetry.instrumentation.starlette import StarletteInstrumentor
 
         # Patch starlette dependencies version check
@@ -28,6 +31,21 @@ def patched_instrumentation_dependencies(self) -> Collection[str]:
         # Apply the patch
         StarletteInstrumentor.instrumentation_dependencies = patched_instrumentation_dependencies
 
+        # Patch to exclude http recieve/send ASGI event spans, this Middleware instrumentation is injected
+        # internally by Starlette Instrumentor, see:
+        # https://github.com/open-telemetry/opentelemetry-python-contrib/blob/51da0a766e5d3cbc746189e10c9573163198cfcd/instrumentation/opentelemetry-instrumentation-asgi/src/opentelemetry/instrumentation/asgi/__init__.py#L573
+        if AGENT_OBSERVABILITY_ENABLED:
+            original_init = OpenTelemetryMiddleware.__init__
+
+            def patched_init(self, app, **kwargs):
+                original_init(self, app, **kwargs)
+                if hasattr(self, "exclude_receive_span"):
+                    self.exclude_receive_span = True
+                if hasattr(self, "exclude_send_span"):
+                    self.exclude_send_span = True
+
+            OpenTelemetryMiddleware.__init__ = patched_init
+
         _logger.debug("Successfully patched Starlette instrumentation_dependencies method")
     except Exception as exc:  # pylint: disable=broad-except
         _logger.warning("Failed to apply Starlette instrumentation patches: %s", exc)

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_starlette_patches.py

@@ -9,19 +9,36 @@
 class TestStarlettePatch(TestCase):
     """Test the Starlette instrumentation patches."""
 
+    @patch("amazon.opentelemetry.distro.patches._starlette_patches.AGENT_OBSERVABILITY_ENABLED", True)
     @patch("amazon.opentelemetry.distro.patches._starlette_patches._logger")
     def test_starlette_patch_applied_successfully(self, mock_logger):
         """Test that the Starlette instrumentation patch is applied successfully."""
         # Create a mock StarletteInstrumentor class
         mock_instrumentor_class = MagicMock()
         mock_instrumentor_class.__name__ = "StarletteInstrumentor"
 
-        # Create a mock module
+        class MockMiddleware:
+            def __init__(self, app, **kwargs):
+                pass
+
+        mock_middleware_class = MockMiddleware
+        original_init = mock_middleware_class.__init__
+
+        # Create mock modules
         mock_starlette_module = MagicMock()
         mock_starlette_module.StarletteInstrumentor = mock_instrumentor_class
 
-        # Mock the import
-        with patch.dict("sys.modules", {"opentelemetry.instrumentation.starlette": mock_starlette_module}):
+        mock_asgi_module = MagicMock()
+        mock_asgi_module.OpenTelemetryMiddleware = mock_middleware_class
+
+        # Mock the imports
+        with patch.dict(
+            "sys.modules",
+            {
+                "opentelemetry.instrumentation.starlette": mock_starlette_module,
+                "opentelemetry.instrumentation.asgi": mock_asgi_module,
+            },
+        ):
             # Apply the patch
             _apply_starlette_instrumentation_patches()
 
@@ -33,6 +50,18 @@ def test_starlette_patch_applied_successfully(self, mock_logger):
             result = mock_instrumentor_class.instrumentation_dependencies(mock_instance)
             self.assertEqual(result, ("starlette >= 0.13",))
 
+            self.assertNotEqual(mock_middleware_class.__init__, original_init)
+
+            # Test middleware patching sets exclude flags
+            mock_middleware_instance = MagicMock()
+            mock_middleware_instance.exclude_receive_span = False
+            mock_middleware_instance.exclude_send_span = False
+
+            mock_middleware_class.__init__(mock_middleware_instance, "app")
+
+            self.assertTrue(mock_middleware_instance.exclude_receive_span)
+            self.assertTrue(mock_middleware_instance.exclude_send_span)
+
             # Verify logging
             mock_logger.debug.assert_called_once_with(
                 "Successfully patched Starlette instrumentation_dependencies method"