ray-project--ray/python/ray/serve/tests/test_tracing_utils.py

import json
import os
import re
import shutil
import sys
import uuid
from pathlib import Path
from threading import Thread
from typing import Set
from unittest.mock import patch

import grpc
import httpx
import pytest
import starlette
from starlette.requests import Request
from starlette.responses import StreamingResponse

import ray
from ray import serve
from ray._common.test_utils import SignalActor
from ray.serve._private.common import ServeComponentType
from ray.serve._private.constants import (
    RAY_SERVE_ENABLE_DIRECT_INGRESS,
    RAY_SERVE_ENABLE_HA_PROXY,
)
from ray.serve._private.logging_utils import get_serve_logs_dir
from ray.serve._private.test_utils import (
    get_application_url,
    wait_for_haproxy_routing_to_replica,
)
from ray.serve._private.tracing_utils import (
    DEFAULT_TRACING_EXPORTER_IMPORT_PATH,
    TRACE_STACK,
    _append_trace_stack,
    _load_span_processors,
    _validate_tracing_exporter,
    _validate_tracing_exporter_processors,
    set_trace_status,
    setup_tracing,
)
from ray.serve.config import HTTPOptions, gRPCOptions
from ray.serve.generated import serve_pb2, serve_pb2_grpc
from ray.serve.grpc_util import gRPCInputStream
from ray.serve.tests.conftest import *  # noqa
from ray.serve.utils import get_trace_context
from ray.tests.conftest import *  # noqa

try:
    from opentelemetry import trace
    from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor
    from opentelemetry.sdk.trace.sampling import ParentBasedTraceIdRatio
    from opentelemetry.trace.propagation.tracecontext import (
        TraceContextTextMapPropagator,
    )
    from opentelemetry.trace.status import StatusCode
except ImportError:
    raise ModuleNotFoundError(
        "`opentelemetry` or `opentelemetry.sdk.trace.export` not found"
    )

CUSTOM_EXPORTER_OUTPUT_FILENAME = "spans.txt"


@pytest.fixture
def use_custom_tracing_exporter():
    yield "ray.serve.tests.test_tracing_utils:custom_tracing_exporter"

    # Clean up output file produced by custom exporter
    if os.path.exists(CUSTOM_EXPORTER_OUTPUT_FILENAME):
        safe_remove(CUSTOM_EXPORTER_OUTPUT_FILENAME)


@pytest.fixture
def serve_and_ray_shutdown():
    serve.shutdown()
    ray.shutdown()
    yield
    serve.shutdown()


class FakeSpan:
    def __init__(self):
        self.status = None

    def set_status(self, status):
        self.status = status


def test_disable_tracing_exporter():
    """Test that setting `tracing_exporter_import_path`
    to an empty string disables tracing.
    """
    is_tracing_setup_successful = setup_tracing(
        component_type=ServeComponentType.REPLICA,
        component_name="component_name",
        component_id="component_id",
        tracing_exporter_import_path="",
        tracing_sampling_ratio=1.0,
    )

    assert is_tracing_setup_successful is False


def test_validate_tracing_exporter_with_string():
    """Test exception message for invalid exporter type."""
    invalid_exporters = [1, "string", []]
    expected_exception = "Tracing exporter must be a function."

    for invalid_exporter in invalid_exporters:
        with pytest.raises(TypeError, match=expected_exception):
            _validate_tracing_exporter(invalid_exporter)


def test_validate_tracing_exporter_with_args():
    """Test exception message for _validate_tracing_exporter.
    if exporter contains an argument.
    """

    def test_exporter(arg):
        return arg

    expected_exception = "Tracing exporter cannot take any arguments."

    with pytest.raises(TypeError, match=expected_exception):
        _validate_tracing_exporter(test_exporter)


def test_validate_tracing_exporter_processors_list():
    """Test exception message for _validate_tracing_exporter.
    if exporter returns invalid return type.
    """
    invalid_span_processors = [1, "string"]
    for invalid_span_processor in invalid_span_processors:
        expected_exception = re.escape(
            "Output of tracing exporter needs to be of type "
            "List[SpanProcessor], but received type "
            f"{type(invalid_span_processor)}."
        )
        with pytest.raises(TypeError, match=expected_exception):
            _validate_tracing_exporter_processors(invalid_span_processor)


def test_validate_tracing_exporter_processors_full_output():
    """Test exception message for _validate_tracing_exporter.
    if exporter returns invalid return type.
    """
    invalid_span_processors = [[1, 2], ["1", "2"]]
    for invalid_span_processor in invalid_span_processors:
        expected_exception = re.escape(
            "Output of tracing exporter needs to be of "
            "type List[SpanProcessor], "
            f"but received type {type(invalid_span_processor[0])}."
        )
        with pytest.raises(TypeError, match=expected_exception):
            _validate_tracing_exporter_processors(invalid_span_processor)


def test_missing_dependencies():
    """Test that setup_tracing raises an exception if
    tracing module is not installed.
    """
    expected_exception = (
        "You must `pip install opentelemetry-api` and "
        "`pip install opentelemetry-sdk` "
        "to enable tracing on Ray Serve."
    )
    with patch("ray.serve._private.tracing_utils.trace", new=None):
        with pytest.raises(ImportError, match=expected_exception):
            setup_tracing(
                component_type=ServeComponentType.REPLICA,
                component_name="component_name",
                component_id="component_id",
                tracing_sampling_ratio=1.0,
            )


def test_default_tracing_exporter(ray_start_cluster):
    """Test that the default tracing exporter returns
    List[SimpleSpanProcessor].
    """
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=1)
    cluster.wait_for_nodes()

    ray.init(address=cluster.address)

    span_processors = _load_span_processors(
        DEFAULT_TRACING_EXPORTER_IMPORT_PATH, "mock_file.json"
    )

    assert isinstance(span_processors, list)

    for span_processor in span_processors:
        assert isinstance(span_processor, SimpleSpanProcessor)


def safe_remove(path):
    """Safely removes a file or directory, handling Windows file locking by using a retry method."""
    if sys.platform != "win32":
        if os.path.isdir(path):
            shutil.rmtree(path)
        else:
            os.remove(path)
        return

    provider = trace.get_tracer_provider()
    if hasattr(provider, "shutdown"):
        provider.shutdown()

    if os.path.isdir(path):
        shutil.rmtree(path)
    else:
        os.remove(path)


def test_custom_tracing_exporter(use_custom_tracing_exporter):
    """Test setup_tracing with a custom tracing exporter."""
    custom_tracing_exporter_path = use_custom_tracing_exporter

    span_processors = _load_span_processors(
        custom_tracing_exporter_path, "mock_file.json"
    )

    assert isinstance(span_processors, list)

    for span_processor in span_processors:
        assert isinstance(span_processor, SimpleSpanProcessor)
        span_processor.shutdown()

    is_tracing_setup_successful = setup_tracing(
        "component_name",
        "component_id",
        ServeComponentType.REPLICA,
        custom_tracing_exporter_path,
        tracing_sampling_ratio=1.0,
    )

    # Validate that tracing is setup successfully
    # and the tracing exporter created the output file
    assert is_tracing_setup_successful
    assert os.path.exists(CUSTOM_EXPORTER_OUTPUT_FILENAME)


def test_tracing_sampler(use_custom_tracing_exporter):
    """Test that tracing sampler is properly configured
    through tracing_sampling_ratio argument.
    """
    custom_tracing_exporter_path = use_custom_tracing_exporter
    tracing_sampling_ratio = 1

    is_tracing_setup_successful = setup_tracing(
        "component_name",
        "component_id",
        ServeComponentType.REPLICA,
        custom_tracing_exporter_path,
        tracing_sampling_ratio,
    )

    # Validate that tracing is setup successfully
    # and the tracing exporter created the output file
    assert is_tracing_setup_successful
    assert os.path.exists(CUSTOM_EXPORTER_OUTPUT_FILENAME)

    tracer = trace.get_tracer(__name__)
    tracer_data = tracer.__dict__

    assert "sampler" in tracer_data

    sampler = tracer_data["sampler"]
    sampler_data = sampler.__dict__

    # ParentBasedTraceIdRatio sampler contains other samplers as attributes
    # The sampling ratio is stored in the underlying samplers
    # Check that we have the expected sampler structure
    assert "_local_parent_not_sampled" in sampler_data
    assert "_local_parent_sampled" in sampler_data
    assert "_remote_parent_sampled" in sampler_data
    assert "_remote_parent_not_sampled" in sampler_data

    assert isinstance(sampler, ParentBasedTraceIdRatio)


@pytest.mark.parametrize(
    (
        "serve_application",
        "expected_proxy_spans_path",
        "expected_replica_spans_path",
        "expected_upstream_spans_path",
    ),
    [
        (
            "basic",
            "fixtures/basic_proxy_spans.json",
            "fixtures/basic_replica_spans.json",
            "fixtures/basic_upstream_spans.json",
        ),
        (
            "streaming",
            "fixtures/streaming_proxy_spans.json",
            "fixtures/streaming_replica_spans.json",
            "fixtures/streaming_upstream_spans.json",
        ),
        (
            "grpc",
            "fixtures/grpc_proxy_spans.json",
            "fixtures/grpc_replica_spans.json",
            "fixtures/grpc_upstream_spans.json",
        ),
    ],
)
def test_tracing_e2e(
    serve_and_ray_shutdown,
    serve_application,
    expected_proxy_spans_path,
    expected_replica_spans_path,
    expected_upstream_spans_path,
):
    """Test tracing e2e."""

    signal_actor = SignalActor.remote()

    @serve.deployment
    class BasicModel:
        def __call__(self, req: starlette.requests.Request):
            replica_context = serve.get_replica_context()
            tracer = trace.get_tracer(__name__)
            with tracer.start_as_current_span(
                "application_span", context=get_trace_context()
            ) as span:
                span.set_attribute("deployment", replica_context.deployment)
                span.set_attribute("replica_id", replica_context.replica_id.unique_id)
                return "hello"

    def hi_gen_sync():
        for i in range(10):
            yield f"hello_{i}"
            # to avoid coalescing chunks
            ray.get(signal_actor.wait.remote())

    @serve.deployment
    class StreamingModel:
        def __call__(self, request: Request) -> StreamingResponse:
            gen = hi_gen_sync()
            replica_context = serve.get_replica_context()
            tracer = trace.get_tracer(__name__)
            with tracer.start_as_current_span(
                "application_span", context=get_trace_context()
            ) as span:
                span.set_attribute("deployment", replica_context.deployment)
                span.set_attribute("replica_id", replica_context.replica_id.unique_id)
                return StreamingResponse(gen, media_type="text/plain")

    @serve.deployment
    class GrpcDeployment:
        def __call__(self, user_message):
            replica_context = serve.get_replica_context()
            tracer = trace.get_tracer(__name__)
            with tracer.start_as_current_span(
                "application_span", context=get_trace_context()
            ) as span:
                span.set_attribute("deployment", replica_context.deployment)
                span.set_attribute("replica_id", replica_context.replica_id.unique_id)

                greeting = f"Hello {user_message.name} from {user_message.foo}"
                num_x2 = user_message.num * 2
                user_response = serve_pb2.UserDefinedResponse(
                    greeting=greeting,
                    num_x2=num_x2,
                )
                return user_response

    if serve_application == "basic":
        serve.start(
            http_options=HTTPOptions(
                host="0.0.0.0",
            )
        )
        serve.run(BasicModel.bind())
        wait_for_haproxy_routing_to_replica()

        setup_tracing(
            component_name="upstream_app",
            component_id="345",
            tracing_sampling_ratio=1.0,
        )
        tracer = trace.get_tracer("test_tracing")
        with tracer.start_as_current_span("upstream_app"):
            ctx = get_trace_context()
            headers = {}
            TraceContextTextMapPropagator().inject(headers, ctx)
            url = get_application_url("HTTP")
            r = httpx.post(f"{url}/", headers=headers)
            assert r.text == "hello"

    elif serve_application == "streaming":
        serve.start(
            http_options=HTTPOptions(
                host="0.0.0.0",
            )
        )
        serve.run(StreamingModel.bind())
        wait_for_haproxy_routing_to_replica()
        setup_tracing(
            component_name="upstream_app",
            component_id="345",
            tracing_sampling_ratio=1.0,
        )
        tracer = trace.get_tracer("test_tracing")
        with tracer.start_as_current_span("upstream_app"):
            ctx = get_trace_context()
            headers = {}
            TraceContextTextMapPropagator().inject(headers, ctx)
            url = get_application_url("HTTP")

            with httpx.stream("GET", f"{url}/", headers=headers) as r:
                r.raise_for_status()
                for i, chunk in enumerate(r.iter_text()):
                    assert chunk == f"hello_{i}"
                    ray.get(signal_actor.send.remote())

    elif serve_application == "grpc":
        # TODO: Remove this once HAProxy supports gRPC
        if RAY_SERVE_ENABLE_HA_PROXY:
            return

        grpc_port = 9000
        grpc_servicer_functions = [
            "ray.serve.generated.serve_pb2_grpc"
            ".add_UserDefinedServiceServicer_to_server",
            "ray.serve.generated.serve_pb2_grpc.add_FruitServiceServicer_to_server",
        ]

        serve.start(
            grpc_options=gRPCOptions(
                port=grpc_port,
                grpc_servicer_functions=grpc_servicer_functions,
            ),
        )
        g = GrpcDeployment.options(name="grpc-deployment").bind()
        serve.run(g)

        setup_tracing(
            component_name="upstream_app",
            component_id="345",
            tracing_sampling_ratio=1.0,
        )
        tracer = trace.get_tracer("test_tracing")
        with tracer.start_as_current_span("upstream_app"):
            ctx = get_trace_context()
            headers = {}
            TraceContextTextMapPropagator().inject(headers, ctx)
            metadata = tuple(headers.items())

            channel = grpc.insecure_channel(get_application_url("gRPC"))

            stub = serve_pb2_grpc.UserDefinedServiceStub(channel)
            request = serve_pb2.UserDefinedMessage(name="foo", num=30, foo="bar")
            response, call = stub.__call__.with_call(request=request, metadata=metadata)

            assert call.code() == grpc.StatusCode.OK, call.code()
            assert response.greeting == "Hello foo from bar", response.greeting

    serve.shutdown()

    serve_logs_dir = get_serve_logs_dir()
    spans_dir = os.path.join(serve_logs_dir, "spans")
    files = os.listdir(spans_dir)

    # proxy, replica, and upstream span files. Under HAProxy the proxy file is
    # written by the fallback ProxyActor on startup but holds no request spans
    # under direct ingress.
    assert len(files) == 3

    replica_filename = None
    proxy_filename = None
    upstream_filename = None
    for file in files:
        if "replica" in file:
            replica_filename = file
        elif "proxy" in file:
            proxy_filename = file
        elif "upstream" in file:
            upstream_filename = file
        else:
            assert False, f"Did not expect tracing file with name {file}"

    assert replica_filename and proxy_filename and upstream_filename

    upstream_spans = load_spans(os.path.join(spans_dir, upstream_filename))
    if not RAY_SERVE_ENABLE_DIRECT_INGRESS:
        proxy_spans = load_spans(os.path.join(spans_dir, proxy_filename))
    else:
        proxy_spans = []
    replica_spans = load_spans(os.path.join(spans_dir, replica_filename))

    entire_trace = replica_spans + proxy_spans + upstream_spans
    validate_span_associations_in_trace(entire_trace)

    expected_upstream_spans = load_json_fixture(expected_upstream_spans_path)
    if not RAY_SERVE_ENABLE_DIRECT_INGRESS:
        expected_proxy_spans = load_json_fixture(expected_proxy_spans_path)
    else:
        expected_proxy_spans = []
    expected_replica_spans = load_json_fixture(expected_replica_spans_path)

    sanitize_spans(upstream_spans)
    sanitize_spans(proxy_spans)
    sanitize_spans(replica_spans)

    assert upstream_spans == expected_upstream_spans
    assert proxy_spans == expected_proxy_spans
    assert replica_spans == expected_replica_spans

    safe_remove(spans_dir)


@pytest.mark.parametrize(
    "protocol,expected_status_code,expected_span_status",
    [
        ("http", 500, StatusCode.ERROR),
        ("streaming", 500, StatusCode.ERROR),
        ("grpc", grpc.StatusCode.INTERNAL.name, StatusCode.ERROR),
    ],
)
def test_tracing_e2e_with_errors(
    serve_and_ray_shutdown, protocol, expected_status_code, expected_span_status
):
    """Test tracing with error responses across HTTP, streaming, and gRPC protocols."""

    @serve.deployment
    class HttpErrorModel:
        def __call__(self, request: starlette.requests.Request):
            replica_context = serve.get_replica_context()
            tracer = trace.get_tracer(__name__)
            with tracer.start_as_current_span(
                "application_span", context=get_trace_context()
            ) as span:
                span.set_attribute("deployment", replica_context.deployment)
                span.set_attribute("replica_id", replica_context.replica_id.unique_id)

                raise RuntimeError("Internal server error")

    @serve.deployment
    class StreamingErrorModel:
        def __call__(self, request: Request) -> StreamingResponse:
            replica_context = serve.get_replica_context()
            tracer = trace.get_tracer(__name__)
            with tracer.start_as_current_span(
                "application_span", context=get_trace_context()
            ) as span:
                span.set_attribute("deployment", replica_context.deployment)
                span.set_attribute("replica_id", replica_context.replica_id.unique_id)

                def error_generator():
                    raise RuntimeError("Streaming error occurred")

                return StreamingResponse(error_generator(), media_type="text/plain")

    @serve.deployment
    class GrpcErrorModel:
        def __call__(self, user_message):
            replica_context = serve.get_replica_context()
            tracer = trace.get_tracer(__name__)
            with tracer.start_as_current_span(
                "application_span", context=get_trace_context()
            ) as span:
                span.set_attribute("deployment", replica_context.deployment)
                span.set_attribute("replica_id", replica_context.replica_id.unique_id)

                # Raise error
                raise RuntimeError("gRPC error occurred")

    # Setup based on protocol
    if protocol == "http":
        serve.start(
            http_options=HTTPOptions(
                host="0.0.0.0",
            )
        )
        serve.run(HttpErrorModel.bind())

        setup_tracing(
            component_name="upstream_app",
            component_id="345",
            tracing_sampling_ratio=1.0,
        )
        tracer = trace.get_tracer("test_tracing")
        with tracer.start_as_current_span("upstream_app"):
            ctx = get_trace_context()
            headers = {}
            TraceContextTextMapPropagator().inject(headers, ctx)

            # Make HTTP request
            url = get_application_url("HTTP")
            response = httpx.post(f"{url}/", headers=headers)
            assert response.status_code == expected_status_code

    elif protocol == "streaming":
        serve.start(
            http_options=HTTPOptions(
                host="0.0.0.0",
            )
        )
        serve.run(StreamingErrorModel.bind())

        setup_tracing(
            component_name="upstream_app",
            component_id="345",
            tracing_sampling_ratio=1.0,
        )
        tracer = trace.get_tracer("test_tracing")
        with tracer.start_as_current_span("upstream_app"):
            ctx = get_trace_context()
            headers = {}
            TraceContextTextMapPropagator().inject(headers, ctx)

            url = get_application_url("HTTP")
            with httpx.stream("GET", f"{url}/", headers=headers) as r:
                assert r.status_code == expected_status_code

    elif protocol == "grpc":
        # TODO: Remove this once HAProxy supports gRPC
        if RAY_SERVE_ENABLE_HA_PROXY:
            return

        grpc_port = 9000
        grpc_servicer_functions = [
            "ray.serve.generated.serve_pb2_grpc.add_UserDefinedServiceServicer_to_server",
        ]

        serve.start(
            grpc_options=gRPCOptions(
                port=grpc_port,
                grpc_servicer_functions=grpc_servicer_functions,
            ),
        )
        serve.run(GrpcErrorModel.options(name="grpc-error-model").bind())

        setup_tracing(
            component_name="upstream_app",
            component_id="345",
            tracing_sampling_ratio=1.0,
        )
        tracer = trace.get_tracer("test_tracing")
        with tracer.start_as_current_span("upstream_app"):
            ctx = get_trace_context()
            headers = {}
            TraceContextTextMapPropagator().inject(headers, ctx)
            channel = grpc.insecure_channel("localhost:9000")
            stub = serve_pb2_grpc.UserDefinedServiceStub(channel)
            request = serve_pb2.UserDefinedMessage(name="test", num=10, foo="bar")

            with pytest.raises(grpc.RpcError) as exception_info:
                _ = stub.__call__(request=request)
            rpc_error = exception_info.value
            print(rpc_error)
            assert rpc_error.code().name == expected_status_code
    else:
        assert False, "Invalid protocol"

    serve.shutdown()

    # Verify the trace data
    serve_logs_dir = get_serve_logs_dir()
    spans_dir = os.path.join(serve_logs_dir, "spans")
    files = os.listdir(spans_dir)

    # proxy, replica, and upstream span files. Under HAProxy the proxy file is
    # written by the fallback ProxyActor on startup but holds no request spans
    # under direct ingress.
    assert len(files) == 3

    replica_filename = None
    proxy_filename = None
    upstream_filename = None
    for file in files:
        if "replica" in file:
            replica_filename = file
        elif "proxy" in file:
            proxy_filename = file
        elif "upstream" in file:
            upstream_filename = file
        else:
            assert False, f"Did not expect tracing file with name {file}"

    assert replica_filename and proxy_filename and upstream_filename

    # Load and check spans
    if not RAY_SERVE_ENABLE_DIRECT_INGRESS:
        proxy_spans = load_spans(os.path.join(spans_dir, proxy_filename))
    else:
        proxy_spans = []
    replica_spans = load_spans(os.path.join(spans_dir, replica_filename))

    # Verify error status in spans
    for span in replica_spans:
        if "application_span" in span["name"]:
            assert span["status"]["status_code"] == expected_span_status.name

            # Check for error attributes based on protocol and error type
            if protocol == "http":
                assert "Internal server error" in str(span.get("events", []))
            elif protocol == "streaming":
                assert "Streaming error occurred" in str(span.get("events", []))
            elif protocol == "grpc":
                assert "gRPC error occurred" in str(span.get("events", []))
            else:
                assert False, "Invalid protocol"

    # Verify status code in proxy spans
    for span in proxy_spans:
        if protocol == "http" or protocol == "streaming":
            if "proxy_http_request" in span["name"]:
                assert (
                    span["attributes"].get("http.status_code") == expected_status_code
                )
                assert span["status"]["status_code"] == "ERROR"
            elif "route_to_replica" in span["name"]:
                pass
            else:
                raise Exception("Invalid proxy span")
        elif protocol == "grpc":
            if "proxy_grpc_request" in span["name"]:
                assert (
                    span["attributes"].get("rpc.grpc.status_code")
                    == expected_status_code
                )
                assert span["status"]["status_code"] == "ERROR"
            elif "route_to_replica" in span["name"]:
                pass
            else:
                assert False, "Invalid proxy span"
        else:
            assert False, "Invalid protocol"
    # Clean up
    safe_remove(spans_dir)


def custom_tracing_exporter():
    """Custom tracing exporter used for testing."""
    out_file = open(CUSTOM_EXPORTER_OUTPUT_FILENAME, "a")

    class FileConsoleSpanExporter(ConsoleSpanExporter):
        def shutdown(self):
            if not out_file.closed:
                out_file.flush()
                out_file.close()

    return [SimpleSpanProcessor(FileConsoleSpanExporter(out=out_file))]


def load_json_fixture(file_path):
    """Load json from a fixture."""
    with Path(__file__).parent.joinpath(file_path).open() as f:
        return json.load(f)


def load_spans(file_path):
    """Load and parse spans from a `.span` file.
    This requires special handling because ConsoleSpanExporter
    does not write proper JSON since the data is streamed.
    """
    if not os.path.exists(file_path) or os.path.getsize(file_path) == 0:
        return []

    with open(file_path, "r") as file:
        file_contents = file.read()

    raw_spans = re.split(r"}\s*\n\s*{", file_contents)
    spans = []
    for i, raw_span in enumerate(raw_spans):
        if len(raw_spans) > 1:
            if i == 0:
                raw_span = raw_span + "}"
            elif i == (len(raw_spans) - 1):
                raw_span = "{" + raw_span
            elif i != 0:
                raw_span = "{" + raw_span + "}"

        span_dict = json.loads(raw_span)

        spans.append(span_dict)

    spans.sort(reverse=True, key=lambda x: x["start_time"])

    return spans


def sanitize_spans(spans):
    """Remove span attributes with ephemeral data."""
    for span in spans:
        del span["context"]
        del span["parent_id"]
        del span["start_time"]
        del span["end_time"]
        del span["resource"]
        for k, _ in span["attributes"].items():
            if "_id" in k:
                span["attributes"][k] = ""


def validate_span_associations_in_trace(spans):
    """Validate that all spans in a
    trace are correctly associated with each
    other through the parent_id relationship.
    """
    if len(spans) <= 1:
        return

    class Span:
        def __init__(self, _span_id: str, _parent_id: str, _name: str):
            self.span_id = _span_id
            self.parent_id = _parent_id
            self.name = _name

        def __repr__(self):
            return f"Span(span_id={self.span_id}, parent_id={self.parent_id}, name={self.name})"

    span_nodes = {}
    starting_span = None
    for span in spans:
        span_id = span["context"]["span_id"].lstrip("0x")
        parent_id = span["parent_id"].lstrip("0x") if span["parent_id"] else None
        name = span["name"]
        new_span = Span(span_id, parent_id, name)
        span_nodes[span_id] = new_span
        # The starting span is always the application span.
        if name == "application_span":
            assert starting_span is None, "Multiple starting spans found in trace"
            starting_span = new_span
    current_span = starting_span
    span_nodes.pop(current_span.span_id)
    while current_span:
        current_span = span_nodes.pop(current_span.parent_id, None)
    # All spans should have been visited.
    assert not span_nodes


def test_set_trace_status_empty_stack():
    """test calling set_trace_status with an empty trace stack.

    When there is nothing in the trace stack, calling set_trace_status does nothing
    and does not error.
    """

    set_trace_status(is_error=True, description="error")
    trace_stack = TRACE_STACK.get([])
    assert trace_stack == []


def test_set_trace_status_error():
    """test calling set_trace_status with error status.

    When there is a trace stack, calling set_trace_status with error updates
    the correct status.
    """
    error_message = "error"
    TRACE_STACK.set([FakeSpan()])
    set_trace_status(is_error=True, description=error_message)
    trace_stack = TRACE_STACK.get([])
    assert len(trace_stack) == 1
    assert trace_stack[0].status.status_code == StatusCode.ERROR
    assert trace_stack[0].status.description == error_message


def test_set_trace_status_ok(caplog):
    """test calling set_trace_status with ok status.

    When there is a trace stack, calling set_trace_status with ok updates
    the correct status.
    """
    ok_message = "ok"
    TRACE_STACK.set([FakeSpan()])
    set_trace_status(is_error=False, description=ok_message)
    trace_stack = TRACE_STACK.get([])
    assert len(trace_stack) == 1
    assert trace_stack[0].status.status_code == StatusCode.OK
    # Note: when the status is OK, the description is not set.
    assert trace_stack[0].status.description is None
    # Ensure we don't attempt to set the description when the status is OK.
    assert (
        "description should only be set when status_code is set to StatusCode.ERROR"
        not in caplog.text
    )


def test_append_trace_stack_multithread():
    """test calling _append_trace_stack in multiple threads.

    When multiple threads call _append_trace_stack, TRACE_STACK should only be updated
    to contain the task name of the thread that called _append_trace_stack.
    """
    number_of_tasks = 10
    passing = set()

    def try_append_trace_stack(_task_name: str, _passing: Set[str]):
        """
        Helper to call `_append_trace_stack()` in a thread. It checks TRACE_STACK in
        the current thread starts out empty and the correct task name is added to it
        after `_append_trace_stack()` is called. If both checks out, the task name will
        be added to the _passing.
        """
        trace_stack_before = TRACE_STACK.get([])
        assert trace_stack_before == []
        _append_trace_stack(_task_name)
        trace_stack_after = TRACE_STACK.get()
        assert trace_stack_after == [_task_name]
        passing.add(_task_name)

    tasks = []
    for i in range(number_of_tasks):
        t = Thread(target=try_append_trace_stack, args=(f"task{i}", passing))
        t.start()
        tasks.append(t)

    for task in tasks:
        task.join()

    assert len(passing) == number_of_tasks


def test_batched_span_attached_to_first_request_trace():
    """Ensure span created inside a batched method attaches to the first request's trace.

    With 6 requests and max_batch_size=3, Serve should create 2 batches. Only the first request
    in each batch should contribute the parent trace to the 'batched_span', yielding exactly
    2 spans whose trace_ids match the traces of those first requests.
    """

    @serve.deployment
    class BatchedDeployment:
        def __init__(self):
            # We keep a counter so we can assert two unique batches producing two unique spans
            self._batch_idx = 0

        @serve.batch(
            max_batch_size=3, batch_wait_timeout_s=1.0, max_concurrent_batches=2
        )
        async def handle_batch(self, reqs):
            tracer = trace.get_tracer(__name__)

            # whichever request is at index 0 is the one whose context is attached
            first_req_id = None
            try:
                first_req_id = reqs[0].headers.get("req-id")
            except Exception:
                pass

            batch_idx = self._batch_idx
            self._batch_idx += 1

            with tracer.start_as_current_span(
                "batched_span", context=get_trace_context()
            ) as span:
                if first_req_id is not None:
                    span.set_attribute("first-req-id", first_req_id)
                span.set_attribute("batch-idx", batch_idx)
                return ["ok" for _ in reqs]

        async def __call__(self, request: Request):
            return await self.handle_batch(request)

    serve.start(http_options=HTTPOptions(host="0.0.0.0"))
    serve.run(BatchedDeployment.bind())

    setup_tracing(
        component_name="upstream_app",
        component_id="batching_test_upstream_multi",
        tracing_sampling_ratio=1.0,
    )

    tracer = trace.get_tracer("test_tracing_batching_multi")

    def do_request(span_name: str):
        req_id = str(uuid.uuid4())
        with tracer.start_as_current_span(span_name) as span:
            # tag the upstream span so we can map req-id -> trace_id later
            span.set_attribute("req-id", req_id)

            # build headers from current OTEL context
            ctx = get_trace_context()
            headers = {"req-id": req_id}
            TraceContextTextMapPropagator().inject(headers, ctx)

            url = get_application_url("HTTP")
            r = httpx.post(f"{url}/", headers=headers)
            r.raise_for_status()

    # Launch 6 requests -> expect 2 batches of 3 requests each
    threads = [Thread(target=do_request, args=(f"upstream_{i}",)) for i in range(6)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    serve.shutdown()

    # Load span files
    serve_logs_dir = get_serve_logs_dir()
    spans_dir = os.path.join(serve_logs_dir, "spans")
    files = os.listdir(spans_dir)
    assert files, "No span files found. Tracing may not be configured."

    replica_filename = None
    upstream_filename = None
    for file in files:
        if "replica" in file:
            replica_filename = file
        elif "upstream" in file:
            upstream_filename = file
    assert replica_filename and upstream_filename

    upstream_spans = load_spans(os.path.join(spans_dir, upstream_filename))
    replica_spans = load_spans(os.path.join(spans_dir, replica_filename))

    id_to_trace = {}
    for s in upstream_spans:
        rid = s.get("attributes", {}).get("req-id")
        if rid:
            id_to_trace[rid] = s["context"]["trace_id"]
    assert (
        len(id_to_trace) == 6
    ), f"Expected 6 upstream request spans, saw {len(id_to_trace)}"

    # Exactly two batch spans, one per each batch, are expected
    batched_spans = [s for s in replica_spans if s["name"] == "batched_span"]
    assert (
        len(batched_spans) == 2
    ), f"Expected 2 batched spans, saw {len(batched_spans)}"

    # For each batched span, its trace_id should equal the trace of the first request in the batch
    batched_trace_ids = set()
    batch_indices = set()
    for bs in batched_spans:
        first_req_id = bs.get("attributes", {}).get("first-req-id")
        assert (
            first_req_id in id_to_trace
        ), f"first-req-id {first_req_id} not found among upstream req-ids"
        assert bs["context"]["trace_id"] == id_to_trace[first_req_id]
        batched_trace_ids.add(bs["context"]["trace_id"])
        batch_indices.add(bs.get("attributes", {}).get("batch-idx"))

    # The two batched spans must correspond to two unique traces
    assert (
        len(batched_trace_ids) == 2
    ), "Batched spans did not map to two unique upstream traces"

    # And they must be from distinct batches
    assert (
        len(batch_indices) == 2
    ), f"Expected two distinct batch indices, got {batch_indices}"

    safe_remove(spans_dir)


@pytest.mark.parametrize(
    "method_name",
    ["ClientStreaming", "BidiStreaming"],
)
def test_grpc_streaming_tracing_attributes(serve_and_ray_shutdown, method_name):
    """Test that tracing attributes are correctly set for gRPC streaming requests.

    Verifies that for client streaming and bidirectional streaming methods,
    the proxy gRPC span contains the expected tracing attributes including
    request_id, deployment, app, request_type, and RPC-specific attributes.
    """
    if RAY_SERVE_ENABLE_HA_PROXY:
        pytest.skip()

    grpc_port = 9000

    serve.start(
        grpc_options=gRPCOptions(
            port=grpc_port,
            grpc_servicer_functions=[
                "ray.serve.generated.serve_pb2_grpc.add_UserDefinedServiceServicer_to_server",
            ],
        ),
    )

    @serve.deployment
    class StreamingDeployment:
        async def ClientStreaming(self, request_stream: gRPCInputStream):
            async for _ in request_stream:
                pass
            return serve_pb2.UserDefinedResponse(greeting="ok", num_x2=0)

        async def BidiStreaming(self, request_stream: gRPCInputStream):
            async for request in request_stream:
                yield serve_pb2.UserDefinedResponse(
                    greeting=f"Hello {request.name}",
                    num_x2=request.num * 2,
                )

    serve.run(StreamingDeployment.options(name="streaming-deployment").bind())

    setup_tracing(
        component_name="upstream_app",
        component_id="345",
        tracing_sampling_ratio=1.0,
    )

    tracer = trace.get_tracer("test_tracing")
    with tracer.start_as_current_span("upstream_streaming"):
        ctx = get_trace_context()
        headers = {}
        TraceContextTextMapPropagator().inject(headers, ctx)
        metadata = tuple(headers.items())

        channel = grpc.insecure_channel("localhost:9000")
        stub = serve_pb2_grpc.UserDefinedServiceStub(channel)

        def request_generator():
            for i in range(2):
                yield serve_pb2.UserDefinedMessage(name=f"msg_{i}", num=i, foo="bar")

        if method_name == "ClientStreaming":
            stub.ClientStreaming(request_generator(), metadata=metadata)
        else:
            list(stub.BidiStreaming(request_generator(), metadata=metadata))

    serve.shutdown()

    # Find proxy spans
    serve_logs_dir = get_serve_logs_dir()
    spans_dir = os.path.join(serve_logs_dir, "spans")
    files = os.listdir(spans_dir)

    proxy_spans = None
    for file in files:
        if "proxy" in file:
            proxy_spans = load_spans(os.path.join(spans_dir, file))
            break

    grpc_proxy_span = next(
        (span for span in proxy_spans if "proxy_grpc_request" in span["name"]),
        None,
    )
    assert grpc_proxy_span is not None
    attrs = grpc_proxy_span["attributes"]

    # Verify tracing attributes are set up correctly
    assert "request_id" in attrs
    assert attrs["deployment"] == "streaming-deployment"
    assert "app" in attrs
    assert attrs["request_type"] == "grpc"
    assert attrs["rpc.system"] == "grpc"
    assert f"/ray.serve.UserDefinedService/{method_name}" == attrs["rpc.method"]
    assert attrs["rpc.grpc.status_code"] == "OK"
    assert grpc_proxy_span["status"]["status_code"] == "OK"

    safe_remove(spans_dir)


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", "-s", __file__]))