skywalking/plugins/sw_grpc.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from asyncio import iscoroutine
from concurrent.futures import Executor
from typing import Any, Awaitable, Callable, List, NamedTuple, Optional, Sequence, Tuple, Union

from skywalking import Component, Layer, config
from skywalking.trace.carrier import Carrier
from skywalking.trace.context import NoopContext, get_context
from skywalking.trace.span import NoopSpan
from skywalking.trace.tags import TagGrpcMethod, TagGrpcStatusCode, TagGrpcUrl

link_vector = ['https://grpc.io/docs/languages/python']
support_matrix = {'grpcio': {'>=3.8': ['1.*']}}
note = """The agent package itself depends on grpcio >= 1.83, which is therefore the
effective minimum version of the instrumented library as well."""


def _get_factory_and_method(rpc_handler: Any) -> Tuple[Callable[..., Any], Callable[..., Any]]:
    import grpc

    if rpc_handler.unary_unary:
        return grpc.unary_unary_rpc_method_handler, rpc_handler.unary_unary
    elif rpc_handler.unary_stream:
        return grpc.unary_stream_rpc_method_handler, rpc_handler.unary_stream
    elif rpc_handler.stream_unary:
        return grpc.stream_unary_rpc_method_handler, rpc_handler.stream_unary
    elif rpc_handler.stream_stream:
        return grpc.stream_stream_rpc_method_handler, rpc_handler.stream_stream
    else:
        raise RuntimeError('RPC handler implementation does not exist')


def _restore_carrier(handler_call_details: Any) -> Carrier:
    metadata_map = dict(handler_call_details.invocation_metadata or ())

    carrier = Carrier()
    for item in carrier:
        val = metadata_map.get(item.key)
        if val is not None:
            item.val = val
    return carrier


def get_method_name(details: Any) -> str:
    method_name = details.method
    if type(method_name) is bytes:
        return method_name.decode()
    if type(method_name) is str:
        return method_name
    return str(method_name)


def install_sync() -> None:
    import grpc
    from grpc import _channel

    def install_server() -> None:
        _grpc_server = grpc.server

        def _make_invoke_intercept_method(
            is_response_streaming: bool,
            next_handler_method: Callable[..., Any],
            handler_call_details: grpc.HandlerCallDetails,
        ) -> Callable[[Any, grpc.ServicerContext], Any]:

            def invoke_intercept_method(request_or_iterator: Any, context: grpc.ServicerContext) -> Any:

                carrier = _restore_carrier(handler_call_details)
                method_name = get_method_name(handler_call_details)
                with (
                    NoopSpan(NoopContext())
                    if config.ignore_grpc_method_check(method_name)
                    else get_context().new_entry_span(op=method_name, carrier=carrier)
                ) as span:
                    span.layer = Layer.RPCFramework
                    span.tag(TagGrpcMethod(method_name))
                    span.component = Component.Grpc
                    span.peer = context.peer()
                    try:
                        return next_handler_method(request_or_iterator, context)
                    except grpc.RpcError as e:
                        if hasattr(e, 'code'):
                            span.tag(TagGrpcStatusCode(e.code()))
                        else:
                            span.tag(TagGrpcStatusCode(grpc.StatusCode.UNKNOWN))
                        span.raised()
                        raise e

            def invoke_intercept_method_response_streaming(
                request_or_iterator: Any, context: grpc.ServicerContext
            ) -> Any:
                carrier = _restore_carrier(handler_call_details)
                method_name = get_method_name(handler_call_details)
                with (
                    NoopSpan(NoopContext())
                    if config.ignore_grpc_method_check(method_name)
                    else get_context().new_entry_span(op=method_name, carrier=carrier)
                ) as span:
                    span.layer = Layer.RPCFramework
                    span.tag(TagGrpcMethod(method_name))
                    span.component = Component.Grpc
                    span.peer = context.peer()
                    try:
                        yield from next_handler_method(request_or_iterator, context)
                    except grpc.RpcError as e:
                        if hasattr(e, 'code'):
                            span.tag(TagGrpcStatusCode(e.code()))
                        else:
                            span.tag(TagGrpcStatusCode(grpc.StatusCode.UNKNOWN))
                        span.raised()
                        raise e

            return invoke_intercept_method_response_streaming if is_response_streaming else invoke_intercept_method

        class ServerInterceptor(grpc.ServerInterceptor):

            def intercept_service(
                self,
                continuation: Callable[[grpc.HandlerCallDetails], grpc.RpcMethodHandler],
                handler_call_details: grpc.HandlerCallDetails,
            ) -> grpc.RpcMethodHandler:
                next_handler = continuation(handler_call_details)

                handler_factory, next_handler_method = _get_factory_and_method(next_handler)

                return handler_factory(
                    _make_invoke_intercept_method(
                        bool(next_handler.response_streaming), next_handler_method, handler_call_details
                    ),
                    request_deserializer=next_handler.request_deserializer,
                    response_serializer=next_handler.response_serializer,
                )

        def _sw_grpc_server(
            thread_pool: Optional[Executor] = None,
            handlers: Optional[Sequence[grpc.GenericRpcHandler]] = None,
            interceptors: Optional[Sequence[Any]] = None,
            options: Optional[Any] = None,
            maximum_concurrent_rpcs: Optional[int] = None,
            compression: Optional[grpc.Compression] = None,
            xds: bool = False,
        ):
            _sw_interceptors = [ServerInterceptor()]
            if interceptors is not None:
                _sw_interceptors.extend(interceptors)
            return _grpc_server(
                thread_pool,
                handlers,
                _sw_interceptors,
                options,
                maximum_concurrent_rpcs,
                compression,
                xds,
            )

        grpc.server = _sw_grpc_server

    def install_client() -> None:
        _grpc_channel = _channel.Channel

        class _ClientCallDetails(NamedTuple):
            method: str
            timeout: Optional[float]
            metadata: Optional[Sequence[Tuple[str, Union[str, bytes]]]]
            credentials: Optional[grpc.CallCredentials]
            wait_for_ready: Optional[bool]
            compression: Any

        class _SWClientCallDetails(_ClientCallDetails, grpc.ClientCallDetails):
            pass

        class _ClientInterceptor(
            grpc.UnaryUnaryClientInterceptor,
            grpc.UnaryStreamClientInterceptor,
            grpc.StreamUnaryClientInterceptor,
            grpc.StreamStreamClientInterceptor,
        ):
            def __init__(self, target: str):
                self.target = target

            def _intercept(
                self,
                continuation: Callable[..., Any],
                client_call_details: grpc.ClientCallDetails,
                request: Any,
            ):
                method_name = get_method_name(client_call_details)
                with (
                    NoopSpan(NoopContext())
                    if config.ignore_grpc_method_check(method_name)
                    else get_context().new_exit_span(op=method_name, peer=self.target, component=Component.Grpc)
                ) as span:
                    span.layer = Layer.RPCFramework
                    span.tag(TagGrpcMethod(method_name))
                    span.tag(TagGrpcUrl(self.target))
                    carrier = span.inject()
                    metadata = list(client_call_details.metadata or [])
                    for item in carrier:
                        metadata.append((item.key, item.val))
                    new_client_call_details = _SWClientCallDetails(
                        client_call_details.method,
                        client_call_details.timeout,
                        metadata,
                        client_call_details.credentials,
                        client_call_details.wait_for_ready,
                        client_call_details.compression,
                    )
                    return continuation(new_client_call_details, request)

            def intercept_unary_unary(self, continuation, client_call_details, request):
                return self._intercept(continuation, client_call_details, request)

            def intercept_unary_stream(self, continuation, client_call_details, request):
                return self._intercept(continuation, client_call_details, request)

            def intercept_stream_unary(self, continuation, client_call_details, request_iterator):
                return self._intercept(continuation, client_call_details, request_iterator)

            def intercept_stream_stream(self, continuation, client_call_details, request_iterator):
                return self._intercept(continuation, client_call_details, request_iterator)

        def _sw_grpc_channel_factory(target: str, *args: Any, **kwargs: Any):
            from skywalking.utils.grpc_channel import is_building_agent_collector_channel

            c = _grpc_channel(target, *args, **kwargs)
            # Prefer agent→OAP build scope: multi-address targets are rewritten (ipv4:/ipv6:)
            # and no longer equal agent_collector_backend_services.
            if (
                is_building_agent_collector_channel()
                or target == config.agent_collector_backend_services
            ):
                return c
            return grpc.intercept_channel(c, _ClientInterceptor(target))

        _channel.Channel = _sw_grpc_channel_factory

    install_client()
    install_server()


def install_async() -> None:
    import grpc
    from grpc.aio import _channel as _aio_channel

    def install_async_server() -> None:
        _grpc_aio_server = grpc.aio.server

        def _make_invoke_intercept_method(
            is_response_streaming: bool,
            next_handler_method: Callable[..., Any],
            handler_call_details: grpc.HandlerCallDetails,
        ) -> Callable[[Any, grpc.aio.ServicerContext[Any, Any]], Awaitable[Any]]:

            async def invoke_intercept_method(
                request_or_iterator: Any, context: grpc.aio.ServicerContext[Any, Any]
            ) -> Any:

                carrier = _restore_carrier(handler_call_details)
                method_name = get_method_name(handler_call_details)
                with (
                    NoopSpan(NoopContext())
                    if config.ignore_grpc_method_check(method_name)
                    else get_context().new_entry_span(op=method_name, carrier=carrier)
                ) as span:
                    span.layer = Layer.RPCFramework
                    span.tag(TagGrpcMethod(method_name))
                    span.component = Component.Grpc
                    span.peer = context.peer()
                    try:
                        return await next_handler_method(request_or_iterator, context)
                    except grpc.RpcError as e:
                        if hasattr(e, 'code'):
                            span.tag(TagGrpcStatusCode(e.code()))
                        else:
                            span.tag(TagGrpcStatusCode(grpc.StatusCode.UNKNOWN))
                        span.raised()
                        raise e

            async def invoke_intercept_method_response_streaming(
                request_or_iterator: Any, context: grpc.aio.ServicerContext[Any, Any]
            ) -> Any:
                carrier = _restore_carrier(handler_call_details)
                method_name = get_method_name(handler_call_details)
                with (
                    NoopSpan(NoopContext())
                    if config.ignore_grpc_method_check(method_name)
                    else get_context().new_entry_span(op=method_name, carrier=carrier)
                ) as span:
                    span.layer = Layer.RPCFramework
                    span.tag(TagGrpcMethod(method_name))
                    span.component = Component.Grpc
                    span.peer = context.peer()
                    try:
                        coroutine_or_asyncgen = next_handler_method(request_or_iterator, context)
                        async for r in (
                            await coroutine_or_asyncgen if iscoroutine(coroutine_or_asyncgen) else coroutine_or_asyncgen
                        ):
                            yield r
                    except grpc.RpcError as e:
                        if hasattr(e, 'code'):
                            span.tag(TagGrpcStatusCode(e.code()))
                        else:
                            span.tag(TagGrpcStatusCode(grpc.StatusCode.UNKNOWN))
                        span.raised()
                        raise e

            return invoke_intercept_method_response_streaming if is_response_streaming else invoke_intercept_method

        class ServerInterceptor(grpc.aio.ServerInterceptor):

            async def intercept_service(
                self,
                continuation: Callable[[grpc.HandlerCallDetails], Awaitable[grpc.RpcMethodHandler]],
                handler_call_details: grpc.HandlerCallDetails,
            ) -> grpc.RpcMethodHandler:
                next_handler = await continuation(handler_call_details)

                handler_factory, next_handler_method = _get_factory_and_method(next_handler)
                return handler_factory(
                    _make_invoke_intercept_method(
                        bool(next_handler.response_streaming), next_handler_method, handler_call_details
                    ),
                    request_deserializer=next_handler.request_deserializer,
                    response_serializer=next_handler.response_serializer,
                )

        def _sw_grpc_aio_server(
            migration_thread_pool: Optional[Executor] = None,
            handlers: Optional[Sequence[grpc.GenericRpcHandler]] = None,
            interceptors: Optional[Sequence[Any]] = None,
            options: Optional[grpc.aio.ChannelArgumentType] = None,
            maximum_concurrent_rpcs: Optional[int] = None,
            compression: Optional[grpc.Compression] = None,
        ):
            _sw_interceptors = [ServerInterceptor()]
            if interceptors is not None:
                _sw_interceptors.extend(interceptors)
            return _grpc_aio_server(
                migration_thread_pool,
                handlers,
                _sw_interceptors,
                options,
                maximum_concurrent_rpcs,
                compression,
            )

        grpc.aio.server = _sw_grpc_aio_server

    def install_async_client() -> None:
        _aio_grpc_channel = _aio_channel.Channel

        class _AioClientCallDetails(NamedTuple):
            method: str
            timeout: Optional[float]
            metadata: Optional[grpc.aio.Metadata]
            credentials: Optional[grpc.CallCredentials]
            wait_for_ready: Optional[bool]

        class _SWAioClientCallDetails(_AioClientCallDetails, grpc.aio.ClientCallDetails):
            pass

        class _AioClientInterceptor:
            def __init__(self, target: str):
                self.target = target

            async def _intercept(
                self,
                continuation: Callable[..., Any],
                client_call_details: grpc.aio.ClientCallDetails,
                request: Any,
            ):
                method_name = get_method_name(client_call_details)
                span = (
                    NoopSpan(NoopContext())
                    if config.ignore_grpc_method_check(method_name)
                    else get_context().new_exit_span(op=method_name, peer=self.target, component=Component.Grpc)
                )
                with span:
                    span.layer = Layer.RPCFramework
                    span.tag(TagGrpcMethod(method_name))
                    span.tag(TagGrpcUrl(self.target))
                    carrier = span.inject()
                    metadata = client_call_details.metadata or grpc.aio.Metadata()
                    for item in carrier:
                        metadata.add(item.key, item.val)
                    new_client_call_details = _SWAioClientCallDetails(
                        client_call_details.method,
                        client_call_details.timeout,
                        metadata,
                        client_call_details.credentials,
                        client_call_details.wait_for_ready,
                    )
                    return await continuation(new_client_call_details, request)

        class _AioClientUnaryUnaryInterceptor(_AioClientInterceptor, grpc.aio.UnaryUnaryClientInterceptor):

            async def intercept_unary_unary(
                self,
                continuation: Callable[
                    [grpc.aio.ClientCallDetails, grpc.aio._typing.RequestType], grpc.aio._call.UnaryUnaryCall
                ],
                client_call_details: grpc.aio.ClientCallDetails,
                request: grpc.aio._typing.RequestType,
            ) -> grpc.aio._call.UnaryUnaryCall:
                return await self._intercept(continuation, client_call_details, request)

        class _AioClientUnaryStreamInterceptor(_AioClientInterceptor, grpc.aio.UnaryStreamClientInterceptor):

            async def intercept_unary_stream(
                self,
                continuation: Callable[
                    [grpc.aio.ClientCallDetails, grpc.aio._typing.RequestType], grpc.aio._call.UnaryStreamCall
                ],
                client_call_details: grpc.aio.ClientCallDetails,
                request: grpc.aio._typing.RequestType,
            ) -> grpc.aio._call.UnaryStreamCall:
                return await self._intercept(continuation, client_call_details, request)

        class _AioClientStreamUnaryInterceptor(_AioClientInterceptor, grpc.aio.StreamUnaryClientInterceptor):

            async def intercept_stream_unary(
                self,
                continuation: Callable[
                    [grpc.aio.ClientCallDetails, grpc.aio._typing.RequestType], grpc.aio._call.StreamUnaryCall
                ],
                client_call_details: grpc.aio.ClientCallDetails,
                request_iterator: grpc.aio._typing.RequestIterableType,
            ) -> grpc.aio._call.StreamUnaryCall:
                return await self._intercept(continuation, client_call_details, request_iterator)

        class _AioClientStreamStreamInterceptor(_AioClientInterceptor, grpc.aio.StreamStreamClientInterceptor):

            async def intercept_stream_stream(
                self,
                continuation: Callable[
                    [grpc.aio.ClientCallDetails, grpc.aio._typing.RequestType], grpc.aio._call.StreamStreamCall
                ],
                client_call_details: grpc.aio.ClientCallDetails,
                request_iterator: grpc.aio._typing.RequestIterableType,
            ) -> grpc.aio._call.StreamStreamCall:
                return await self._intercept(continuation, client_call_details, request_iterator)

        # Must remain a class assigned to the module-global name `Channel`:
        # since grpcio 1.83.0, grpc.aio._channel._BaseMultiCallable.__init__ validates
        # `isinstance(self._references[0], Channel)` against this module global, so
        # rebinding it to a plain factory function raises
        # `TypeError: isinstance() arg 2 must be a type, ...` on every stub creation.
        class _SWAioChannel(_aio_grpc_channel):
            def __init__(
                self,
                target: str,
                options: grpc.aio.ChannelArgumentType,
                credentials: Optional[grpc.ChannelCredentials],
                compression: Optional[grpc.Compression],
                interceptors: Optional[Sequence[grpc.aio.ClientInterceptor]],
            ):
                from skywalking.utils.grpc_channel import is_building_agent_collector_channel

                # Multi-address collector targets are rewritten; do not rely on string equality alone.
                skip_sw = (
                    is_building_agent_collector_channel()
                    or target == config.agent_collector_backend_services
                )
                if not skip_sw:
                    _sw_interceptors: List[grpc.aio.ClientInterceptor] = [
                        _AioClientUnaryUnaryInterceptor(target),
                        _AioClientUnaryStreamInterceptor(target),
                        _AioClientStreamUnaryInterceptor(target),
                        _AioClientStreamStreamInterceptor(target),
                    ]
                    _sw_interceptors.extend(interceptors or [])
                    interceptors = _sw_interceptors
                super().__init__(target, options, credentials, compression, interceptors)

        _aio_channel.Channel = _SWAioChannel

    install_async_client()
    install_async_server()


def install():
    install_sync()
    install_async()