codekingpro/portable-devtools

1import random2 3from eventlet import wsgi4from eventlet.zipkin import api5from eventlet.zipkin._thrift.zipkinCore.constants import \6    SERVER_RECV, SERVER_SEND7from eventlet.zipkin.http import \8    HDR_TRACE_ID, HDR_SPAN_ID, HDR_PARENT_SPAN_ID, HDR_SAMPLED9 10 11_sampler = None12__original_handle_one_response__ = wsgi.HttpProtocol.handle_one_response13 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 server22        trace_id = span_id = api.generate_trace_id()23        parent_id = None24        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_rate46 47    def sampling(self):48        # avoid generating unneeded random numbers49        if self.sampling_rate == 1.0:50            return True51        r = random.random()52        if r < self.sampling_rate:53            return True54        return False55 56 57def int_or_none(val):58    if val is None:59        return None60    return int(val, 16)61 62 63def bool_or_none(val):64    if val == '1':65        return True66    if val == '0':67        return False68    return None69 70 71def patch(sampling_rate):72    global _sampler73    _sampler = Sampler(sampling_rate)74    wsgi.HttpProtocol.handle_one_response = _patched_handle_one_response75 76 77def unpatch():78    wsgi.HttpProtocol.handle_one_response = __original_handle_one_response__79