From 556dc2e4d6a1ba792494383f0030ac4f8871ee40 Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Wed, 3 Apr 2024 18:02:50 +0000 Subject: [PATCH] [Core] Add new attribute to resent HTTP request spans `DistributedTracingPolicy` will now set an attribute, `http.request.resend_count`, on HTTP spans for resent requests to indicate the resend attempt count. Signed-off-by: Paul Van Eck --- sdk/core/azure-core/CHANGELOG.md | 2 + .../pipeline/policies/_distributed_tracing.py | 3 + .../azure/core/pipeline/policies/_retry.py | 1 + .../core/pipeline/policies/_retry_async.py | 1 + .../async_tests/test_tracing_policy_async.py | 58 +++++++++++++++++++ .../azure-core/tests/test_tracing_policy.py | 43 +++++++++++++- 6 files changed, 106 insertions(+), 2 deletions(-) create mode 100644 sdk/core/azure-core/tests/async_tests/test_tracing_policy_async.py diff --git a/sdk/core/azure-core/CHANGELOG.md b/sdk/core/azure-core/CHANGELOG.md index ed3148ec9656..4f7193a81a4a 100644 --- a/sdk/core/azure-core/CHANGELOG.md +++ b/sdk/core/azure-core/CHANGELOG.md @@ -4,6 +4,8 @@ ### Features Added +- Tracing: `DistributedTracingPolicy` will now set an attribute, `http.request.resend_count`, on HTTP spans for resent requests to indicate the resend attempt number. #35069 + ### Breaking Changes ### Bugs Fixed diff --git a/sdk/core/azure-core/azure/core/pipeline/policies/_distributed_tracing.py b/sdk/core/azure-core/azure/core/pipeline/policies/_distributed_tracing.py index 6a7619eb1cbf..003052664e49 100644 --- a/sdk/core/azure-core/azure/core/pipeline/policies/_distributed_tracing.py +++ b/sdk/core/azure-core/azure/core/pipeline/policies/_distributed_tracing.py @@ -76,6 +76,7 @@ class DistributedTracingPolicy(SansIOHTTPPolicy[HTTPRequestType, HTTPResponseTyp TRACING_CONTEXT = "TRACING_CONTEXT" _REQUEST_ID = "x-ms-client-request-id" _RESPONSE_ID = "x-ms-request-id" + _HTTP_RESEND_COUNT = "http.request.resend_count" def __init__(self, **kwargs: Any): self._network_span_namer = kwargs.get("network_span_namer", _default_network_span_namer) @@ -125,6 +126,8 @@ def end_span( http_request: Union[HttpRequest, LegacyHttpRequest] = request.http_request if span is not None: span.set_http_attributes(http_request, response=response) + if request.context.get("retry_count"): + span.add_attribute(self._HTTP_RESEND_COUNT, request.context["retry_count"]) request_id = http_request.headers.get(self._REQUEST_ID) if request_id is not None: span.add_attribute(self._REQUEST_ID, request_id) diff --git a/sdk/core/azure-core/azure/core/pipeline/policies/_retry.py b/sdk/core/azure-core/azure/core/pipeline/policies/_retry.py index 12e2a1985894..2fd4d7b987ce 100644 --- a/sdk/core/azure-core/azure/core/pipeline/policies/_retry.py +++ b/sdk/core/azure-core/azure/core/pipeline/policies/_retry.py @@ -528,6 +528,7 @@ def send(self, request: PipelineRequest[HTTPRequestType]) -> PipelineResponse[HT ) try: self._configure_timeout(request, absolute_timeout, is_response_error) + request.context["retry_count"] = len(retry_settings["history"]) response = self.next.send(request) if self.is_retry(retry_settings, response): retry_active = self.increment(retry_settings, response=response) diff --git a/sdk/core/azure-core/azure/core/pipeline/policies/_retry_async.py b/sdk/core/azure-core/azure/core/pipeline/policies/_retry_async.py index 1e63e3b09fe7..a61228e4f739 100644 --- a/sdk/core/azure-core/azure/core/pipeline/policies/_retry_async.py +++ b/sdk/core/azure-core/azure/core/pipeline/policies/_retry_async.py @@ -176,6 +176,7 @@ async def send( ) try: self._configure_timeout(request, absolute_timeout, is_response_error) + request.context["retry_count"] = len(retry_settings["history"]) response = await self.next.send(request) if self.is_retry(retry_settings, response): retry_active = self.increment(retry_settings, response=response) diff --git a/sdk/core/azure-core/tests/async_tests/test_tracing_policy_async.py b/sdk/core/azure-core/tests/async_tests/test_tracing_policy_async.py new file mode 100644 index 000000000000..355e3499bc36 --- /dev/null +++ b/sdk/core/azure-core/tests/async_tests/test_tracing_policy_async.py @@ -0,0 +1,58 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +"""Tests for the distributed tracing policy in an async pipeline.""" + +import pytest + + +from azure.core.pipeline import AsyncPipeline +from azure.core.pipeline.policies import AsyncRetryPolicy, DistributedTracingPolicy +from azure.core.pipeline.transport import ( + HttpResponse, + AsyncHttpTransport, +) +from azure.core.settings import settings + +from tracing_common import FakeSpan +from utils import HTTP_REQUESTS + + +@pytest.mark.asyncio +@pytest.mark.parametrize("http_request", HTTP_REQUESTS) +async def test_span_retry_attributes(http_request): + class MockTransport(AsyncHttpTransport): + def __init__(self): + self._count = 0 + + async def __aexit__(self, exc_type, exc_val, exc_tb): + pass + + async def close(self): + pass + + async def open(self): + pass + + async def send(self, request, **kwargs): + self._count += 1 + response = HttpResponse(request, None) + response.status_code = 429 + return response + + http_request = http_request("GET", "http://localhost/") + retry_policy = AsyncRetryPolicy(retry_total=2) + distributed_tracing_policy = DistributedTracingPolicy() + transport = MockTransport() + + settings.tracing_implementation.set_value(FakeSpan) + with FakeSpan(name="parent") as root_span: + pipeline = AsyncPipeline(transport, [retry_policy, distributed_tracing_policy]) + await pipeline.run(http_request) + + assert transport._count == 3 + assert len(root_span.children) == 3 + assert root_span.children[0].attributes.get("http.request.resend_count") is None + assert root_span.children[1].attributes.get("http.request.resend_count") == 1 + assert root_span.children[2].attributes.get("http.request.resend_count") == 2 diff --git a/sdk/core/azure-core/tests/test_tracing_policy.py b/sdk/core/azure-core/tests/test_tracing_policy.py index c89ff86471d5..425b00bc94dc 100644 --- a/sdk/core/azure-core/tests/test_tracing_policy.py +++ b/sdk/core/azure-core/tests/test_tracing_policy.py @@ -5,8 +5,9 @@ """Tests for the distributed tracing policy.""" import logging -from azure.core.pipeline import PipelineResponse, PipelineRequest, PipelineContext -from azure.core.pipeline.policies import DistributedTracingPolicy, UserAgentPolicy +from azure.core.pipeline import Pipeline, PipelineResponse, PipelineRequest, PipelineContext +from azure.core.pipeline.policies import DistributedTracingPolicy, UserAgentPolicy, RetryPolicy +from azure.core.pipeline.transport import HttpTransport from azure.core.settings import settings from tracing_common import FakeSpan import time @@ -210,6 +211,44 @@ def test_distributed_tracing_policy_with_user_agent(http_request, http_response) assert network_span.status == "Transport trouble" +@pytest.mark.parametrize("http_request,http_response", request_and_responses_product(HTTP_RESPONSES)) +def test_span_retry_attributes(http_request, http_response): + class MockTransport(HttpTransport): + def __init__(self): + self._count = 0 + + def __exit__(self, exc_type, exc_val, exc_tb): + pass + + def close(self): + pass + + def open(self): + pass + + def send(self, request, **kwargs): + self._count += 1 + response = create_http_response(http_response, request, None) + response.status_code = 429 + return response + + settings.tracing_implementation.set_value(FakeSpan) + + http_request = http_request("GET", "http://localhost/") + retry_policy = RetryPolicy(retry_total=2) + distributed_tracing_policy = DistributedTracingPolicy() + transport = MockTransport() + + with FakeSpan(name="parent") as root_span: + pipeline = Pipeline(transport, [retry_policy, distributed_tracing_policy]) + pipeline.run(http_request) + assert transport._count == 3 + assert len(root_span.children) == 3 + assert root_span.children[0].attributes.get("http.request.resend_count") is None + assert root_span.children[1].attributes.get("http.request.resend_count") == 1 + assert root_span.children[2].attributes.get("http.request.resend_count") == 2 + + @pytest.mark.parametrize("http_request,http_response", request_and_responses_product(HTTP_RESPONSES)) def test_span_namer(http_request, http_response): settings.tracing_implementation.set_value(FakeSpan)