grpc_channelz.py

Go to the documentation of this file.

# Copyright 2020 gRPC authors.
#
# Licensed 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.
"""
This contains helpers for gRPC services defined in
https://github.com/grpc/grpc-proto/blob/master/grpc/channelz/v1/channelz.proto
"""
import ipaddress
import logging
from typing import Iterator, Optional
 
import grpc
from grpc_channelz.v1 import channelz_pb2
from grpc_channelz.v1 import channelz_pb2_grpc
 
import framework.rpc
 
logger = logging.getLogger(__name__)
 
# Type aliases
# Channel
Channel = channelz_pb2.Channel
ChannelConnectivityState = channelz_pb2.ChannelConnectivityState
ChannelState = ChannelConnectivityState.State  # pylint: disable=no-member
_GetTopChannelsRequest = channelz_pb2.GetTopChannelsRequest
_GetTopChannelsResponse = channelz_pb2.GetTopChannelsResponse
# Subchannel
Subchannel = channelz_pb2.Subchannel
_GetSubchannelRequest = channelz_pb2.GetSubchannelRequest
_GetSubchannelResponse = channelz_pb2.GetSubchannelResponse
# Server
Server = channelz_pb2.Server
_GetServersRequest = channelz_pb2.GetServersRequest
_GetServersResponse = channelz_pb2.GetServersResponse
# Sockets
Socket = channelz_pb2.Socket
SocketRef = channelz_pb2.SocketRef
_GetSocketRequest = channelz_pb2.GetSocketRequest
_GetSocketResponse = channelz_pb2.GetSocketResponse
Address = channelz_pb2.Address
Security = channelz_pb2.Security
# Server Sockets
_GetServerSocketsRequest = channelz_pb2.GetServerSocketsRequest
_GetServerSocketsResponse = channelz_pb2.GetServerSocketsResponse
 
 
class ChannelzServiceClient(framework.rpc.grpc.GrpcClientHelper):
    stub: channelz_pb2_grpc.ChannelzStub
 
    def __init__(self, channel: grpc.Channel):
        super().__init__(channel, channelz_pb2_grpc.ChannelzStub)
 
    @staticmethod
    def is_sock_tcpip_address(address: Address):
        return address.WhichOneof('address') == 'tcpip_address'
 
    @staticmethod
    def is_ipv4(tcpip_address: Address.TcpIpAddress):
        # According to proto, tcpip_address.ip_address is either IPv4 or IPv6.
        # Correspondingly, it's either 4 bytes or 16 bytes in length.
        return len(tcpip_address.ip_address) == 4
 
    @classmethod
    def sock_address_to_str(cls, address: Address):
        if cls.is_sock_tcpip_address(address):
            tcpip_address: Address.TcpIpAddress = address.tcpip_address
            if cls.is_ipv4(tcpip_address):
                ip = ipaddress.IPv4Address(tcpip_address.ip_address)
            else:
                ip = ipaddress.IPv6Address(tcpip_address.ip_address)
            return f'{ip}:{tcpip_address.port}'
        else:
            raise NotImplementedError('Only tcpip_address implemented')
 
    @classmethod
    def sock_addresses_pretty(cls, socket: Socket):
        return (f'local={cls.sock_address_to_str(socket.local)}, '
                f'remote={cls.sock_address_to_str(socket.remote)}')
 
    @staticmethod
    def find_server_socket_matching_client(server_sockets: Iterator[Socket],
                                           client_socket: Socket) -> Socket:
        for server_socket in server_sockets:
            if server_socket.remote == client_socket.local:
                return server_socket
        return None
 
    def find_channels_for_target(self, target: str,
                                 **kwargs) -> Iterator[Channel]:
        return (channel for channel in self.list_channels(**kwargs)
                if channel.data.target == target)
 
    def find_server_listening_on_port(self, port: int,
                                      **kwargs) -> Optional[Server]:
        for server in self.list_servers(**kwargs):
            listen_socket_ref: SocketRef
            for listen_socket_ref in server.listen_socket:
                listen_socket = self.get_socket(listen_socket_ref.socket_id,
                                                **kwargs)
                listen_address: Address = listen_socket.local
                if (self.is_sock_tcpip_address(listen_address) and
                        listen_address.tcpip_address.port == port):
                    return server
        return None
 
    def list_channels(self, **kwargs) -> Iterator[Channel]:
        """
        Iterate over all pages of all root channels.
 
        Root channels are those which application has directly created.
        This does not include subchannels nor non-top level channels.
        """
        start: int = -1
        response: Optional[_GetTopChannelsResponse] = None
        while start < 0 or not response.end:
            # From proto: To request subsequent pages, the client generates this
            # value by adding 1 to the highest seen result ID.
            start += 1
            response = self.call_unary_with_deadline(
                rpc='GetTopChannels',
                req=_GetTopChannelsRequest(start_channel_id=start),
                **kwargs)
            for channel in response.channel:
                start = max(start, channel.ref.channel_id)
                yield channel
 
    def list_servers(self, **kwargs) -> Iterator[Server]:
        """Iterate over all pages of all servers that exist in the process."""
        start: int = -1
        response: Optional[_GetServersResponse] = None
        while start < 0 or not response.end:
            # From proto: To request subsequent pages, the client generates this
            # value by adding 1 to the highest seen result ID.
            start += 1
            response = self.call_unary_with_deadline(
                rpc='GetServers',
                req=_GetServersRequest(start_server_id=start),
                **kwargs)
            for server in response.server:
                start = max(start, server.ref.server_id)
                yield server
 
    def list_server_sockets(self, server: Server, **kwargs) -> Iterator[Socket]:
        """List all server sockets that exist in server process.
 
        Iterating over the results will resolve additional pages automatically.
        """
        start: int = -1
        response: Optional[_GetServerSocketsResponse] = None
        while start < 0 or not response.end:
            # From proto: To request subsequent pages, the client generates this
            # value by adding 1 to the highest seen result ID.
            start += 1
            response = self.call_unary_with_deadline(
                rpc='GetServerSockets',
                req=_GetServerSocketsRequest(server_id=server.ref.server_id,
                                             start_socket_id=start),
                **kwargs)
            socket_ref: SocketRef
            for socket_ref in response.socket_ref:
                start = max(start, socket_ref.socket_id)
                # Yield actual socket
                yield self.get_socket(socket_ref.socket_id, **kwargs)
 
    def list_channel_sockets(self, channel: Channel,
                             **kwargs) -> Iterator[Socket]:
        """List all sockets of all subchannels of a given channel."""
        for subchannel in self.list_channel_subchannels(channel, **kwargs):
            yield from self.list_subchannels_sockets(subchannel, **kwargs)
 
    def list_channel_subchannels(self, channel: Channel,
                                 **kwargs) -> Iterator[Subchannel]:
        """List all subchannels of a given channel."""
        for subchannel_ref in channel.subchannel_ref:
            yield self.get_subchannel(subchannel_ref.subchannel_id, **kwargs)
 
    def list_subchannels_sockets(self, subchannel: Subchannel,
                                 **kwargs) -> Iterator[Socket]:
        """List all sockets of a given subchannel."""
        for socket_ref in subchannel.socket_ref:
            yield self.get_socket(socket_ref.socket_id, **kwargs)
 
    def get_subchannel(self, subchannel_id, **kwargs) -> Subchannel:
        """Return a single Subchannel, otherwise raises RpcError."""
        response: _GetSubchannelResponse = self.call_unary_with_deadline(
            rpc='GetSubchannel',
            req=_GetSubchannelRequest(subchannel_id=subchannel_id),
            **kwargs)
        return response.subchannel
 
    def get_socket(self, socket_id, **kwargs) -> Socket:
        """Return a single Socket, otherwise raises RpcError."""
        response: _GetSocketResponse = self.call_unary_with_deadline(
            rpc='GetSocket',
            req=_GetSocketRequest(socket_id=socket_id),
            **kwargs)
        return response.socket