Spaces:
Sleeping
Sleeping
File size: 2,336 Bytes
6b4889b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
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)
|