Upload 2 files

aworld/trace/baggage/sofa_tracer.py

@@ -0,0 +1,128 @@
+from aworld.trace.base import Propagator, Carrier, TraceContext
+from aworld.trace.baggage import BaggageContext
+from aworld.logs.util import logger
+from aworld.trace.base import AttributeValueType
+
+
+class SofaTracerBaggagePropagator(Propagator):
+    """
+    Sofa tracer baggage propagator.
+    """
+
+    _TRACE_ID_HEDER_NAMES = ["SOFA-TraceId", "sofaTraceId"]
+    _SPAN_ID_HEDER_NAMES = ["SOFA-RpcId", "sofaRpcId"]
+    _PEN_ATTRS_HEDER_NAME = "sofaPenAttrs"
+    _SYS_PEN_ATTRS_HEDER_NAME = "sysPenAttrs"
+
+    _TRACE_ID_BAGGAGE_KEY = "attributes.sofa.traceid"
+    _SPAN_ID_BAGGAGE_KEY = "attributes.sofa.rpcid"
+    _PEN_ATTRS_BAGGAGE_KEY = "attributes.sofa.penattrs"
+    _SYS_PEN_ATTRS_BAGGAGE_KEY = "attributes.sofa.syspenattrs"
+
+    def extract(self, carrier: Carrier):
+        """
+        Extract trace context from carrier.
+        Args:
+            carrier: The carrier to extract trace context from.
+        Returns:
+            A dict of trace context.
+        """
+        trace_id = None
+        span_id = None
+        for name in self._TRACE_ID_HEDER_NAMES:
+            trace_id = self._get_value(carrier, name)
+            if trace_id:
+                break
+        for name in self._SPAN_ID_HEDER_NAMES:
+            span_id = self._get_value(carrier, name)
+            if span_id:
+                break
+        pen_attrs = self._get_value(carrier, self._PEN_ATTRS_HEDER_NAME)
+        sys_pen_attrs = self._get_value(
+            carrier, self._SYS_PEN_ATTRS_HEDER_NAME)
+
+        logger.info(
+            f"extract trace_id: {trace_id}, span_id: {span_id}, pen_attrs: {pen_attrs}, sys_pen_attrs: {sys_pen_attrs}")
+        if trace_id and span_id:
+            BaggageContext.set_baggage(self._TRACE_ID_BAGGAGE_KEY, trace_id)
+            span_id = span_id + ".1"
+            BaggageContext.set_baggage(self._SPAN_ID_BAGGAGE_KEY, span_id)
+            if pen_attrs:
+                BaggageContext.set_baggage(
+                    self._PEN_ATTRS_BAGGAGE_KEY, pen_attrs)
+            if sys_pen_attrs:
+                BaggageContext.set_baggage(
+                    self._SYS_PEN_ATTRS_BAGGAGE_KEY, sys_pen_attrs)
+
+    def inject(self, trace_context: TraceContext, carrier: Carrier):
+        """
+        Inject trace context to carrier.
+        Args:
+            trace_context: The trace context to inject.
+            carrier: The carrier to inject trace context to.
+        """
+        baggage = BaggageContext.get_baggage()
+
+        if baggage:
+            trace_id = baggage.get(self._TRACE_ID_BAGGAGE_KEY)
+            span_id = baggage.get(self._SPAN_ID_BAGGAGE_KEY)
+            if trace_id and span_id:
+                carrier.set(self._TRACE_ID_HEDER_NAMES[0], trace_id)
+                carrier.set(self._SPAN_ID_HEDER_NAMES[0], span_id)
+
+            pen_attrs_dict = {}
+            for key, value in baggage.items():
+                if key == self._TRACE_ID_BAGGAGE_KEY or key == self._SPAN_ID_BAGGAGE_KEY:
+                    continue
+                if key == self._PEN_ATTRS_BAGGAGE_KEY and value:
+                    pen_attrs_dict.update(dict(item.split("=")
+                                          for item in value.split("&")))
+                    continue
+                if key == self._SYS_PEN_ATTRS_BAGGAGE_KEY and value:
+                    carrier.set(self._SYS_PEN_ATTRS_HEDER_NAME, value)
+                    continue
+
+                # other baggage items will be injected to sofaPenAttrs
+                pen_attrs_dict.update({key: value})
+
+            if pen_attrs_dict:
+                pen_attrs = "&".join(f"{key}={value}"
+                                     for key, value in pen_attrs_dict.items())
+                carrier.set(self._PEN_ATTRS_HEDER_NAME, pen_attrs)
+
+
+class SofaSpanHelper:
+    """
+    Sofa span helper.
+    """
+
+    @staticmethod
+    def set_sofa_context_to_attr(span_attributes: dict[str, AttributeValueType]):
+        """
+        Set sofa context to span attributes.
+        Args:
+            span_attributes: The span attributes to set sofa context to.
+        """
+        baggage = BaggageContext.get_baggage()
+        if baggage:
+            trace_id = baggage.get(
+                SofaTracerBaggagePropagator._TRACE_ID_BAGGAGE_KEY)
+            span_id = baggage.get(
+                SofaTracerBaggagePropagator._SPAN_ID_BAGGAGE_KEY)
+            if trace_id and span_id:
+                span_attributes.update({
+                    SofaTracerBaggagePropagator._TRACE_ID_BAGGAGE_KEY: trace_id,
+                    SofaTracerBaggagePropagator._SPAN_ID_BAGGAGE_KEY: span_id
+                })
+            pen_attrs = baggage.get(
+                SofaTracerBaggagePropagator._PEN_ATTRS_BAGGAGE_KEY)
+            if pen_attrs:
+                span_attributes.update({
+                    SofaTracerBaggagePropagator._PEN_ATTRS_BAGGAGE_KEY: pen_attrs
+                })
+            sys_pen_attrs = baggage.get(
+                SofaTracerBaggagePropagator._SYS_PEN_ATTRS_BAGGAGE_KEY)
+            if sys_pen_attrs:
+                span_attributes.update({
+                    SofaTracerBaggagePropagator._SYS_PEN_ATTRS_BAGGAGE_KEY: sys_pen_attrs
+                })

aworld/trace/baggage/w3c.py

@@ -0,0 +1,65 @@
+import re
+from typing import List
+from aworld.trace.base import Propagator, Carrier, TraceContext
+from aworld.trace.baggage import BaggageContext
+from aworld.logs.util import logger
+from urllib.parse import quote_plus, unquote_plus
+
+
+class W3CBaggagePropagator(Propagator):
+    """
+    W3C baggage propagator.
+    """
+
+    _MAX_HEADER_LENGTH = 8192
+    _MAX_PAIR_LENGTH = 4096
+    _MAX_PAIRS = 180
+    _BAGGAGE_HEADER_NAME = "baggage"
+    _DELIMITER_PATTERN = re.compile(r"[ \t]*,[ \t]*")
+
+    def extract(self, carrier: Carrier):
+        """
+        Extract the trace context from the carrier.
+        Args:
+            carrier: The carrier to extract the trace context from.
+        """
+        baggage_header = self._get_value(carrier, self._BAGGAGE_HEADER_NAME)
+        if not baggage_header:
+            return None
+
+        if len(baggage_header) > self._MAX_HEADER_LENGTH:
+            logger.warning(
+                f"baggage header length exceeds {self._MAX_HEADER_LENGTH}")
+            return None
+
+        baggage_entries: List[str] = re.split(
+            self._DELIMITER_PATTERN, baggage_header)
+        if len(baggage_entries) > self._MAX_PAIRS:
+            logger.warning(f"baggage entries exceeds {self._MAX_PAIRS}")
+
+        for entry in baggage_entries:
+            if len(entry) > self._MAX_PAIR_LENGTH:
+                logger.warning(
+                    f"baggage entry length exceeds {self._MAX_PAIR_LENGTH}")
+                continue
+            try:
+                key, value = entry.split("=", 1)
+                key = unquote_plus(key).strip()
+                value = unquote_plus(value).strip()
+            except ValueError:
+                logger.warning(f"baggage entry format error: {entry}")
+                continue
+            BaggageContext.set_baggage(key, value)
+
+    def inject(self, carrier: Carrier, context: TraceContext):
+        """
+        Inject the trace context into the carrier.
+        Args:
+            carrier: The carrier to inject the trace context into.
+            context: The trace context to inject.
+        """
+        baggage = BaggageContext.get_baggage()
+        if baggage:
+            baggage_header = ",".join(
+                f"{quote_plus(key)}={quote_plus(value)}" for key, value in baggage.items())
+            carrier.set(self._BAGGAGE_HEADER_NAME, baggage_header)