1import random 2 3from eventlet import wsgi 4from eventlet.zipkin import api 5from eventlet.zipkin._thrift.zipkinCore.constants import \ 6 SERVER_RECV, SERVER_SEND 7from eventlet.zipkin.http import \ 8 HDR_TRACE_ID, HDR_SPAN_ID, HDR_PARENT_SPAN_ID, HDR_SAMPLED 9 10 11_sampler = None 12__original_handle_one_response__ = wsgi.HttpProtocol.handle_one_response 13 14 15def _patched_handle_one_response(self): 16 api.init_trace_data() 17 trace_id = int_or_none(self.headers.getheader(HDR_TRACE_ID)) 18 span_id = int_or_none(self.headers.getheader(HDR_SPAN_ID)) 19 parent_id = int_or_none(self.headers.getheader(HDR_PARENT_SPAN_ID)) 20 sampled = bool_or_none(self.headers.getheader(HDR_SAMPLED)) 21 if trace_id is None: # front-end server 22 trace_id = span_id = api.generate_trace_id() 23 parent_id = None 24 sampled = _sampler.sampling() 25 ip, port = self.request.getsockname()[:2] 26 ep = api.ZipkinDataBuilder.build_endpoint(ip, port) 27 trace_data = api.TraceData(name=self.command, 28 trace_id=trace_id, 29 span_id=span_id, 30 parent_id=parent_id, 31 sampled=sampled, 32 endpoint=ep) 33 api.set_trace_data(trace_data) 34 api.put_annotation(SERVER_RECV) 35 api.put_key_value('http.uri', self.path) 36 37 __original_handle_one_response__(self) 38 39 if api.is_sample(): 40 api.put_annotation(SERVER_SEND) 41 42 43class Sampler(object): 44 def __init__(self, sampling_rate): 45 self.sampling_rate = sampling_rate 46 47 def sampling(self): 48 # avoid generating unneeded random numbers 49 if self.sampling_rate == 1.0: 50 return True 51 r = random.random() 52 if r < self.sampling_rate: 53 return True 54 return False 55 56 57def int_or_none(val): 58 if val is None: 59 return None 60 return int(val, 16) 61 62 63def bool_or_none(val): 64 if val == '1': 65 return True 66 if val == '0': 67 return False 68 return None 69 70 71def patch(sampling_rate): 72 global _sampler 73 _sampler = Sampler(sampling_rate) 74 wsgi.HttpProtocol.handle_one_response = _patched_handle_one_response 75 76 77def unpatch(): 78 wsgi.HttpProtocol.handle_one_response = __original_handle_one_response__ 79