aws_advanced_python_wrapper/utils/telemetry/open_telemetry.py

#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
#  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.

from __future__ import annotations

from typing import TYPE_CHECKING, Callable, Optional, Sequence, Union

if TYPE_CHECKING:
    from opentelemetry.util.types import AttributeValue

from opentelemetry import context as context_api
from opentelemetry import trace
from opentelemetry import trace as trace_api
from opentelemetry.metrics import (CallbackOptions, Meter, Observation,
                                   get_meter)
from opentelemetry.sdk.trace import ReadableSpan, Span, StatusCode, Tracer

from aws_advanced_python_wrapper.utils.log import Logger
from aws_advanced_python_wrapper.utils.messages import Messages
from aws_advanced_python_wrapper.utils.telemetry.telemetry import (
    TelemetryConst, TelemetryContext, TelemetryCounter, TelemetryFactory,
    TelemetryGauge, TelemetryTraceLevel)

logger = Logger(__name__)

INSTRUMENTATION_NAME = "aws-advanced-python-wrapper"


class OpenTelemetryContext(TelemetryContext):
    def __init__(self, tracer: Tracer, name: str, trace_level: TelemetryTraceLevel,
                 start_time: Optional[int] = None, link_span: Optional[Span] = None):
        self._name = name
        self._tracer = tracer
        self._span: Optional[Span]
        self._meter: Meter
        self._token: Optional[object] = None

        current_span: Span = trace.get_current_span()  # type: ignore
        is_root = (current_span is None or current_span == trace.INVALID_SPAN)

        if is_root and trace_level == TelemetryTraceLevel.NESTED:
            trace_level = TelemetryTraceLevel.TOP_LEVEL

        links: Sequence[trace_api.Link] = ()
        if trace_level in [TelemetryTraceLevel.FORCE_TOP_LEVEL, TelemetryTraceLevel.TOP_LEVEL]:
            if link_span is not None:
                links = [trace.Link(link_span.get_span_context())]
            else:
                if not is_root:
                    links = [trace.Link(current_span.get_span_context())]

            self._span = self._tracer.start_span(self._name, context=context_api.Context(),
                                                 links=links, start_time=start_time)  # type: ignore
            if not is_root:
                self.set_attribute(TelemetryConst.TRACE_NAME_ANNOTATION, self._name)

            ctx = trace.set_span_in_context(self._span)  # type: ignore
            self._token = context_api.attach(ctx)
            logger.debug("OpenTelemetryContext.TelemetryTraceID", self._name,
                         self._span.get_span_context().trace_id)  # type: ignore

        elif trace_level == TelemetryTraceLevel.NESTED:
            if link_span is not None:
                links = [trace.Link(link_span.get_span_context())]

            self._span = self._tracer.start_span(self._name, links=links, start_time=start_time)  # type: ignore
            ctx = trace.set_span_in_context(self._span)  # type: ignore
            self._token = context_api.attach(ctx)
            self.set_attribute(TelemetryConst.TRACE_NAME_ANNOTATION, self._name)

        elif trace_level == TelemetryTraceLevel.NO_TRACE:
            self._span = None

    def set_success(self, success: bool):
        if self._span is not None:
            self._span.set_status(StatusCode.OK if success else StatusCode.ERROR)

    def set_attribute(self, key: str, value: AttributeValue):
        if self._span is not None:
            self._span.set_attribute(key, value)

    def set_exception(self, exception: Exception):
        if self._span is not None and exception is not None:
            self._span.set_attribute(TelemetryConst.EXCEPTION_TYPE_ANNOTATION, exception.__class__.__name__)
            self._span.set_attribute(TelemetryConst.EXCEPTION_MESSAGE_ANNOTATION, str(exception))
            self._span.record_exception(exception)

    def get_name(self) -> str:
        return self._name

    def close_context(self):
        if self._token is not None:
            context_api.detach(self._token)
        if self._span is not None:
            self._span.end()

    @property
    def tracer(self) -> Tracer:
        return self._tracer

    @property
    def span(self) -> Optional[Span]:
        return self._span


def post_copy(context: OpenTelemetryContext, trace_level: TelemetryTraceLevel):
    if trace_level == TelemetryTraceLevel.NO_TRACE:
        return

    _clone_and_close_context(context, trace_level)


def _clone_and_close_context(context: OpenTelemetryContext, trace_level: TelemetryTraceLevel) -> OpenTelemetryContext:
    if not isinstance(context.span, ReadableSpan):
        raise RuntimeError(Messages.get("OpenTelemetry.InvalidContext"))

    clone = OpenTelemetryContext(
        context.tracer, TelemetryConst.COPY_TRACE_NAME_PREFIX + context.get_name(),
        trace_level, context.span.start_time, context.span)

    for key in context.span.attributes:  # type: ignore
        value = context.span.attributes[key]  # type: ignore
        clone.set_attribute(key, value)

    clone.span.set_status(context.span.status)  # type: ignore
    clone.set_attribute(TelemetryConst.SOURCE_TRACE_ANNOTATION, str(context.span.get_span_context().trace_id))
    clone.span.end(context.span.end_time)  # type: ignore

    return clone


class OpenTelemetryCounter(TelemetryCounter):
    def __init__(self, meter: Meter, name: str):
        self._meter: Meter = meter
        self._name: str = name
        self._counter = meter.create_up_down_counter(self._name, unit="1")

    def add(self, value):
        self._counter.add(value)

    def inc(self):
        self._counter.add(1)

    def get_name(self):
        return self._name


class OpenTelemetryGauge(TelemetryGauge):
    def __init__(self, meter: Meter, name: str, callback: Callable[[], Union[float, int]]):
        self._meter: Meter = meter
        self.name: str = name
        self._counter = meter.create_observable_up_down_counter(name, callbacks=[self._callback_observation], unit="1")
        self._callback: Callable[[], float] = callback

    def get_name(self):
        return self.name

    def _callback_observation(self, options: CallbackOptions):
        value: Union[int, float] = self._callback()
        observation = Observation(value)
        yield observation


class OpenTelemetryFactory(TelemetryFactory):
    def open_telemetry_context(self, name: str, trace_level: TelemetryTraceLevel) -> TelemetryContext:
        return OpenTelemetryContext(trace.get_tracer(INSTRUMENTATION_NAME), name, trace_level)  # type: ignore

    def post_copy(self, context: TelemetryContext, trace_level: TelemetryTraceLevel):
        if isinstance(context, OpenTelemetryContext):
            post_copy(context, trace_level)
        else:
            raise RuntimeError(Messages.get_formatted("OpenTelemetryFactory.WrongParameterType", type(context)))

    def create_counter(self, name: str) -> TelemetryCounter:
        return OpenTelemetryCounter(get_meter(INSTRUMENTATION_NAME), name)

    def create_gauge(self, name: str, callback: Callable[[], Union[float, int]]):
        return OpenTelemetryGauge(get_meter(INSTRUMENTATION_NAME), name, callback)