azext_edge/edge/providers/stats.py

# coding=utf-8
# ----------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License file in the project root for license information.
# ----------------------------------------------------------------------------------------------

import binascii
import json
from datetime import datetime
from typing import TYPE_CHECKING, List, Optional, Tuple, Union

from azure.cli.core.azclierror import ResourceNotFoundError
from knack.log import get_logger
from rich.console import Console

from ..common import AIO_BROKER_DIAGNOSTICS_SERVICE, PROTOBUF_SERVICE_API_PORT, PodState
from ..util import get_timestamp_now_utc
from .base import V1Pod, get_namespaced_pods_by_prefix, portforward_socket

logger = get_logger(__name__)

console = Console(highlight=True)

if TYPE_CHECKING:
    # pylint: disable=no-name-in-module
    from socket import socket
    from zipfile import ZipInfo

    from opentelemetry.proto.trace.v1.trace_pb2 import TracesData


def _preprocess_stats(
    namespace: Optional[str] = None, diag_service_pod_prefix: str = AIO_BROKER_DIAGNOSTICS_SERVICE
) -> Tuple[str, V1Pod]:
    if not namespace:
        from .base import DEFAULT_NAMESPACE

        namespace = DEFAULT_NAMESPACE

    target_pods = get_namespaced_pods_by_prefix(prefix=diag_service_pod_prefix, namespace=namespace)
    if not target_pods:
        raise ResourceNotFoundError(
            f"Diagnostics service pod '{diag_service_pod_prefix}' does not exist in namespace '{namespace}'."
        )
    for pod in target_pods:
        if pod.status.phase.lower() == PodState.running.value:
            return namespace, pod

    raise ResourceNotFoundError(
        f"No diagnostics service pod '{diag_service_pod_prefix}' in phase "
        f"'{PodState.running.value}' detected in namespace '{namespace}'."
    )


def get_traces(
    namespace: Optional[str] = None,
    diag_service_pod_prefix: str = AIO_BROKER_DIAGNOSTICS_SERVICE,
    pod_protobuf_port: int = PROTOBUF_SERVICE_API_PORT,
    trace_ids: Optional[List[str]] = None,
    trace_dir: Optional[str] = None,
) -> Union[List["TracesData"], List[Tuple["ZipInfo", str]], None]:
    """
    trace_ids: List[str] hex representation of trace Ids.
    """
    if not any([trace_ids, trace_dir]):
        raise ValueError("At least trace_ids or trace_dir is required.")

    from zipfile import ZIP_DEFLATED, ZipFile, ZipInfo

    from google.protobuf.json_format import MessageToDict
    from rich.progress import MofNCompleteColumn, Progress

    from ..util import normalize_dir

    # pylint: disable=no-name-in-module
    from .proto.diagnostics_service_pb2 import Request, Response, TraceRetrievalInfo

    namespace, diagnostic_pod = _preprocess_stats(namespace=namespace, diag_service_pod_prefix=diag_service_pod_prefix)

    for_support_bundle = False
    trace_ids = trace_ids or []

    if trace_ids:
        if trace_ids[0] == "!support_bundle!":
            trace_ids.pop()
            for_support_bundle = True
        trace_ids = [binascii.unhexlify(t) for t in trace_ids]

    with Progress(
        *Progress.get_default_columns(),
        MofNCompleteColumn(),
        transient=False,
        disable=bool(trace_ids) or for_support_bundle,
    ) as progress:
        with portforward_socket(
            namespace=namespace, pod_name=diagnostic_pod.metadata.name, pod_port=pod_protobuf_port
        ) as socket:
            request = Request(get_traces=TraceRetrievalInfo(trace_ids=trace_ids))
            serialized_request = request.SerializeToString()
            request_len_b = len(serialized_request).to_bytes(4, byteorder="big")

            socket.sendall(request_len_b)
            socket.sendall(serialized_request)

            traces: List[dict] = []
            if trace_dir:
                normalized_dir_path = normalize_dir(dir_path=trace_dir)
                normalized_dir_path = normalized_dir_path.joinpath(
                    f"broker_traces_{get_timestamp_now_utc(format='%Y%m%dT%H%M%S')}.zip"
                )
                # pylint: disable=consider-using-with
                myzip = ZipFile(file=str(normalized_dir_path), mode="w", compression=ZIP_DEFLATED)

            progress_set = False
            progress_task = None
            total_trace_count = 0
            current_trace_count = 0
            try:
                while True:
                    if current_trace_count and current_trace_count >= total_trace_count:
                        break
                    rbytes = _fetch_bytes(socket, 4)
                    response_size = int.from_bytes(rbytes, byteorder="big")
                    response_bytes = _fetch_bytes(socket, response_size)

                    if response_bytes == b"":
                        logger.warning("TCP socket closed. Trace processing aborted.")
                        return

                    response = Response.FromString(response_bytes)
                    current_trace_count = current_trace_count + 1

                    if not total_trace_count:
                        total_trace_count = response.retrieved_trace.total_trace_count
                        if total_trace_count == 0:
                            logger.warning("No traces to fetch. Processing aborted.")
                            break

                    if not progress.disable and not progress_set:
                        progress_task = progress.add_task(
                            "[deep_sky_blue4]Gathering traces...", total=response.retrieved_trace.total_trace_count
                        )
                        progress_set = True

                    msg_dict = MessageToDict(message=response.retrieved_trace.trace, use_integers_for_enums=True)
                    root_span, resource_name, timestamp = _determine_root_span(message_dict=msg_dict)

                    if progress_set:
                        progress.update(progress_task, advance=1)
                    if not all([root_span, resource_name, timestamp]):
                        logger.debug("Could not process root span. Skipping trace.")
                        continue

                    span_trace_id = root_span["traceId"]
                    span_name = root_span["name"]

                    if trace_ids:
                        traces.append(msg_dict)
                    if trace_dir or for_support_bundle:
                        archive = f"{resource_name}.{span_name}.{span_trace_id}"
                        pb_suffix = ".otlp.pb"
                        tempo_suffix = ".tempo.json"

                        datetime_tuple = tuple(timestamp.timetuple())
                        zinfo_pb = ZipInfo(filename=f"{archive}{pb_suffix}", date_time=datetime_tuple)
                        # Fixed in Py 3.9 https://github.com/python/cpython/issues/70373
                        zinfo_pb.file_size = 0
                        zinfo_pb.compress_size = 0

                        zinfo_tempo = ZipInfo(filename=f"{archive}{tempo_suffix}", date_time=datetime_tuple)
                        zinfo_tempo.file_size = 0
                        zinfo_tempo.compress_size = 0

                        otlp_format_pair = (zinfo_pb, response.retrieved_trace.trace.SerializeToString())
                        tempo_format_pair = (
                            zinfo_tempo,
                            json.dumps(_convert_otlp_to_tempo(msg_dict), sort_keys=True),
                        )

                        if for_support_bundle:
                            traces.append(otlp_format_pair)
                            traces.append(tempo_format_pair)
                            continue

                        # Original OTLP
                        myzip.writestr(
                            zinfo_or_arcname=otlp_format_pair[0],
                            data=otlp_format_pair[1],
                        )
                        # Tempo
                        myzip.writestr(
                            zinfo_or_arcname=tempo_format_pair[0],
                            data=tempo_format_pair[1],
                        )

                if traces:
                    return traces

            finally:
                if trace_dir:
                    myzip.close()


def _determine_root_span(message_dict: dict) -> Tuple[str, str, Union[datetime, None]]:
    """
    Attempts to determine root span, and normalizes traceId, spanId and parentSpanId to hex.
    """
    import base64

    root_span = None
    resource_name = None
    timestamp = None

    for resource_span in message_dict.get("resourceSpans", []):
        for scope_span in resource_span.get("scopeSpans", []):
            for span in scope_span.get("spans", []):
                if "traceId" in span:
                    span["traceId"] = base64.b64decode(span["traceId"]).hex()
                if "spanId" in span:
                    span["spanId"] = base64.b64decode(span["spanId"]).hex()
                if "parentSpanId" in span:
                    span["parentSpanId"] = base64.b64decode(span["parentSpanId"]).hex()
                else:
                    root_span = span

                    if "startTimeUnixNano" in root_span:
                        timestamp_unix_nano = root_span["startTimeUnixNano"]
                        timestamp = datetime.utcfromtimestamp(float(timestamp_unix_nano) / 1e9)

                    # determine resource name
                    resource = resource_span.get("resource", {})
                    attributes = resource.get("attributes", [])
                    for a in attributes:
                        if a["key"] == "service.name":
                            resource_name = a["value"].get("stringValue", "unknown")

    return root_span, resource_name, timestamp


def _convert_otlp_to_tempo(message_dict: dict) -> dict:
    """
    Convert OTLP payload to Grafana Tempo.
    """
    from copy import deepcopy

    new_dict = deepcopy(message_dict)

    new_dict["batches"] = new_dict.pop("resourceSpans")
    for batch in new_dict.get("batches", []):
        batch["instrumentationLibrarySpans"] = batch.pop("scopeSpans", [])
        for inst_lib_span in batch.get("instrumentationLibrarySpans", []):
            inst_lib_span["instrumentationLibrary"] = inst_lib_span.pop("scope", {})

    return new_dict


def _fetch_bytes(socket: "socket", size: int) -> bytes:
    result_bytes = socket.recv(size)
    if result_bytes == b"":
        return result_bytes

    result_bytes_len = len(result_bytes)
    while result_bytes_len < size:
        remaining_bytes_size = size - result_bytes_len
        interm_bytes = socket.recv(remaining_bytes_size)
        if interm_bytes == b"":
            break
        result_bytes += interm_bytes
        result_bytes_len += len(interm_bytes)

    return result_bytes