import hashlib import logging import random import socket import threading import uuid from collections import deque from collections.abc import Sequence from datetime import datetime from typing import Final from urllib.parse import urljoin import httpx from opentelemetry import trace as trace_api from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.sdk.util.instrumentation import InstrumentationScope from opentelemetry.semconv._incubating.attributes.deployment_attributes import ( # type: ignore[import-untyped] DEPLOYMENT_ENVIRONMENT, ) from opentelemetry.semconv._incubating.attributes.host_attributes import ( # type: ignore[import-untyped] HOST_NAME, ) from opentelemetry.semconv.attributes import service_attributes from opentelemetry.trace import Link, SpanContext, TraceFlags from configs import dify_config from dify_trace_aliyun.entities.aliyun_trace_entity import SpanData from dify_trace_aliyun.entities.semconv import ACS_ARMS_SERVICE_FEATURE INVALID_SPAN_ID: Final[int] = 0x0000000000000000 INVALID_TRACE_ID: Final[int] = 0x00000000000000000000000000000000 DEFAULT_TIMEOUT: Final[int] = 5 DEFAULT_MAX_QUEUE_SIZE: Final[int] = 1000 DEFAULT_SCHEDULE_DELAY_SEC: Final[int] = 5 DEFAULT_MAX_EXPORT_BATCH_SIZE: Final[int] = 50 logger = logging.getLogger(__name__) class TraceClient: def __init__( self, service_name: str, endpoint: str, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE, schedule_delay_sec: int = DEFAULT_SCHEDULE_DELAY_SEC, max_export_batch_size: int = DEFAULT_MAX_EXPORT_BATCH_SIZE, ): self.endpoint = endpoint self.resource = Resource( attributes={ service_attributes.SERVICE_NAME: service_name, service_attributes.SERVICE_VERSION: f"dify-{dify_config.project.version}-{dify_config.COMMIT_SHA}", DEPLOYMENT_ENVIRONMENT: f"{dify_config.DEPLOY_ENV}-{dify_config.EDITION}", HOST_NAME: socket.gethostname(), ACS_ARMS_SERVICE_FEATURE: "genai_app", } ) self.span_builder = SpanBuilder(self.resource) self.exporter = OTLPSpanExporter(endpoint=endpoint) self.max_queue_size = max_queue_size self.schedule_delay_sec = schedule_delay_sec self.max_export_batch_size = max_export_batch_size self.queue: deque = deque(maxlen=max_queue_size) self.condition = threading.Condition(threading.Lock()) self.done = False self.worker_thread = threading.Thread(target=self._worker, daemon=True) self.worker_thread.start() self._spans_dropped = False def export(self, spans: Sequence[ReadableSpan]): self.exporter.export(spans) def api_check(self) -> bool: try: response = httpx.head(self.endpoint, timeout=DEFAULT_TIMEOUT) if response.status_code == 405: return True else: logger.warning("AliyunTrace API check failed: Unexpected status code: %s", response.status_code) return False except httpx.RequestError as e: logger.warning("AliyunTrace API check failed: %s", str(e)) raise ValueError(f"AliyunTrace API check failed: {str(e)}") def get_project_url(self) -> str: return "https://arms.console.aliyun.com/#/llm" def add_span(self, span_data: SpanData | None) -> None: if span_data is None: return span: ReadableSpan = self.span_builder.build_span(span_data) with self.condition: if len(self.queue) == self.max_queue_size: if not self._spans_dropped: logger.warning("Queue is full, likely spans will be dropped.") self._spans_dropped = True self.queue.appendleft(span) if len(self.queue) >= self.max_export_batch_size: self.condition.notify() def _worker(self) -> None: while not self.done: with self.condition: if len(self.queue) < self.max_export_batch_size and not self.done: self.condition.wait(timeout=self.schedule_delay_sec) self._export_batch() def _export_batch(self) -> None: spans_to_export: list[ReadableSpan] = [] with self.condition: while len(spans_to_export) < self.max_export_batch_size and self.queue: spans_to_export.append(self.queue.pop()) if spans_to_export: try: self.exporter.export(spans_to_export) except Exception as e: logger.warning("Error exporting spans: %s", e) def shutdown(self) -> None: with self.condition: self.done = True self.condition.notify_all() self.worker_thread.join() self._export_batch() self.exporter.shutdown() class SpanBuilder: def __init__(self, resource: Resource) -> None: self.resource = resource self.instrumentation_scope = InstrumentationScope( __name__, "", None, None, ) def build_span(self, span_data: SpanData) -> ReadableSpan: span_context = trace_api.SpanContext( trace_id=span_data.trace_id, span_id=span_data.span_id, is_remote=False, trace_flags=trace_api.TraceFlags(trace_api.TraceFlags.SAMPLED), trace_state=None, ) parent_span_context = None if span_data.parent_span_id is not None: parent_span_context = trace_api.SpanContext( trace_id=span_data.trace_id, span_id=span_data.parent_span_id, is_remote=False, trace_flags=trace_api.TraceFlags(trace_api.TraceFlags.SAMPLED), trace_state=None, ) span = ReadableSpan( name=span_data.name, context=span_context, parent=parent_span_context, resource=self.resource, attributes=span_data.attributes, events=span_data.events, links=span_data.links, kind=span_data.span_kind, status=span_data.status, start_time=span_data.start_time, end_time=span_data.end_time, instrumentation_scope=self.instrumentation_scope, ) return span def create_link(trace_id_str: str) -> Link: placeholder_span_id = INVALID_SPAN_ID try: trace_id = int(trace_id_str, 16) except ValueError as e: raise ValueError(f"Invalid trace ID format: {trace_id_str}") from e span_context = SpanContext( trace_id=trace_id, span_id=placeholder_span_id, is_remote=False, trace_flags=TraceFlags(TraceFlags.SAMPLED) ) return Link(span_context) def generate_span_id() -> int: span_id = random.getrandbits(64) while span_id == INVALID_SPAN_ID: span_id = random.getrandbits(64) return span_id def convert_to_trace_id(uuid_v4: str | None) -> int: if uuid_v4 is None: raise ValueError("UUID cannot be None") try: uuid_obj = uuid.UUID(uuid_v4) return uuid_obj.int except ValueError as e: raise ValueError(f"Invalid UUID input: {uuid_v4}") from e def convert_string_to_id(string: str | None) -> int: if not string: return generate_span_id() hash_bytes = hashlib.sha256(string.encode("utf-8")).digest() return int.from_bytes(hash_bytes[:8], byteorder="big", signed=False) def convert_to_span_id(uuid_v4: str | None, span_type: str) -> int: if uuid_v4 is None: raise ValueError("UUID cannot be None") try: uuid_obj = uuid.UUID(uuid_v4) except ValueError as e: raise ValueError(f"Invalid UUID input: {uuid_v4}") from e combined_key = f"{uuid_obj.hex}-{span_type}" return convert_string_to_id(combined_key) def convert_datetime_to_nanoseconds(start_time_a: datetime | None) -> int | None: if start_time_a is None: return None timestamp_in_seconds = start_time_a.timestamp() return int(timestamp_in_seconds * 1e9) def build_endpoint(base_url: str, license_key: str) -> str: if "log.aliyuncs.com" in base_url: # cms2.0 endpoint return urljoin(base_url, f"adapt_{license_key}/api/v1/traces") else: # xtrace endpoint return urljoin(base_url, f"adapt_{license_key}/api/otlp/traces")