impuls/lib/python3.11/site-packages/eventlet/zipkin/wsgi.py

79 lines
2.2 KiB
Python

import random
from eventlet import wsgi
from eventlet.zipkin import api
from eventlet.zipkin._thrift.zipkinCore.constants import \
SERVER_RECV, SERVER_SEND
from eventlet.zipkin.http import \
HDR_TRACE_ID, HDR_SPAN_ID, HDR_PARENT_SPAN_ID, HDR_SAMPLED
_sampler = None
__original_handle_one_response__ = wsgi.HttpProtocol.handle_one_response
def _patched_handle_one_response(self):
api.init_trace_data()
trace_id = int_or_none(self.headers.getheader(HDR_TRACE_ID))
span_id = int_or_none(self.headers.getheader(HDR_SPAN_ID))
parent_id = int_or_none(self.headers.getheader(HDR_PARENT_SPAN_ID))
sampled = bool_or_none(self.headers.getheader(HDR_SAMPLED))
if trace_id is None: # front-end server
trace_id = span_id = api.generate_trace_id()
parent_id = None
sampled = _sampler.sampling()
ip, port = self.request.getsockname()[:2]
ep = api.ZipkinDataBuilder.build_endpoint(ip, port)
trace_data = api.TraceData(name=self.command,
trace_id=trace_id,
span_id=span_id,
parent_id=parent_id,
sampled=sampled,
endpoint=ep)
api.set_trace_data(trace_data)
api.put_annotation(SERVER_RECV)
api.put_key_value('http.uri', self.path)
__original_handle_one_response__(self)
if api.is_sample():
api.put_annotation(SERVER_SEND)
class Sampler(object):
def __init__(self, sampling_rate):
self.sampling_rate = sampling_rate
def sampling(self):
# avoid generating unneeded random numbers
if self.sampling_rate == 1.0:
return True
r = random.random()
if r < self.sampling_rate:
return True
return False
def int_or_none(val):
if val is None:
return None
return int(val, 16)
def bool_or_none(val):
if val == '1':
return True
if val == '0':
return False
return None
def patch(sampling_rate):
global _sampler
_sampler = Sampler(sampling_rate)
wsgi.HttpProtocol.handle_one_response = _patched_handle_one_response
def unpatch():
wsgi.HttpProtocol.handle_one_response = __original_handle_one_response__