trace.py (12811B)
1 #Embedded file name: C:\BuildAgent\work\20d73f1b24eed06e\eve\release\V24.01\packages\beeline\trace.py 2 import base64 3 import copy 4 import datetime 5 import functools 6 import hashlib 7 import json 8 import math 9 import random 10 import struct 11 import sys 12 import threading 13 import traceback 14 import inspect 15 from collections import defaultdict 16 from contextlib import contextmanager 17 from beeline.internal import log, stringify_exception 18 import beeline.propagation 19 import beeline.propagation.honeycomb 20 MAX_INT32 = math.pow(2, 32) - 1 21 SPAN_ID_BYTES = 8 22 TRACE_ID_BYTES = 16 23 24 class Trace(object): 25 26 def __init__(self, trace_id, dataset = None): 27 self.id = trace_id 28 self.dataset = dataset 29 self.stack = [] 30 self.fields = {} 31 self.rollup_fields = defaultdict(float) 32 33 def copy(self): 34 result = Trace(self.id) 35 result.stack = copy.copy(self.stack) 36 result.fields = copy.copy(self.fields) 37 return result 38 39 40 class Tracer(object): 41 42 def __init__(self, client): 43 self._client = client 44 self.presend_hook = None 45 self.sampler_hook = None 46 self.http_trace_parser_hook = beeline.propagation.honeycomb.http_trace_parser_hook 47 self.http_trace_propagation_hook = beeline.propagation.honeycomb.http_trace_propagation_hook 48 49 @contextmanager 50 def __call__(self, name, trace_id = None, parent_id = None): 51 try: 52 span = None 53 if self.get_active_trace_id() and trace_id is None: 54 span = self.start_span(context={'name': name}, parent_id=parent_id) 55 if span: 56 log('tracer context manager started new span, id = %s', span.id) 57 else: 58 span = self.start_trace(context={'name': name}, trace_id=trace_id, parent_span_id=parent_id) 59 if span: 60 log('tracer context manager started new trace, id = %s', span.trace_id) 61 yield span 62 except Exception as e: 63 if span: 64 span.add_context({'app.exception_type': str(type(e)), 65 'app.exception_string': stringify_exception(e), 66 'app.exception_stacktrace': traceback.format_exc()}) 67 raise 68 finally: 69 if span: 70 if span.is_root(): 71 log('tracer context manager ending trace, id = %s', span.trace_id) 72 self.finish_trace(span) 73 else: 74 log('tracer context manager ending span, id = %s', span.id) 75 self.finish_span(span) 76 else: 77 log('tracer context manager span for %s was unexpectedly None', name) 78 79 def start_trace(self, context = None, trace_id = None, parent_span_id = None, dataset = None): 80 if trace_id: 81 if self._trace: 82 log('warning: start_trace got explicit trace_id but we are already in a trace. starting new trace with id = %s', trace_id) 83 self._trace = Trace(trace_id, dataset) 84 else: 85 self._trace = Trace(generate_trace_id(), dataset) 86 return self.start_span(context=context, parent_id=parent_span_id) 87 88 def start_span(self, context = None, parent_id = None): 89 if not self._trace: 90 log('start_span called but no trace is active') 91 return 92 span_id = generate_span_id() 93 if parent_id: 94 parent_span_id = parent_id 95 else: 96 parent_span_id = self._trace.stack[-1].id if self._trace.stack else None 97 ev = self._client.new_event(data=self._trace.fields) 98 if context: 99 ev.add(data=context) 100 ev.add(data={'trace.trace_id': self._trace.id, 101 'trace.parent_id': parent_span_id, 102 'trace.span_id': span_id}) 103 is_root = len(self._trace.stack) == 0 104 span = Span(trace_id=self._trace.id, parent_id=parent_span_id, id=span_id, event=ev, is_root=is_root) 105 self._trace.stack.append(span) 106 return span 107 108 def finish_span(self, span): 109 if span is None: 110 return 111 if span.event: 112 if self._trace: 113 if self._trace.dataset: 114 span.event.dataset = self._trace.dataset 115 if span.is_root(): 116 for k, v in self._trace.rollup_fields.items(): 117 span.event.add_field(k, v) 118 119 for k, v in span.rollup_fields.items(): 120 span.event.add_field(k, v) 121 122 for k, v in self._trace.fields.items(): 123 if k not in span.event.fields(): 124 span.event.add_field(k, v) 125 126 duration = datetime.datetime.now() - span.event.start_time 127 duration_ms = duration.total_seconds() * 1000.0 128 span.event.add_field('duration_ms', duration_ms) 129 self._run_hooks_and_send(span) 130 else: 131 log('warning: span has no event, was it initialized correctly?') 132 if not self._trace: 133 log('warning: span finished without an active trace') 134 return 135 if span.trace_id != self._trace.id: 136 log('warning: finished span called for span in inactive trace. current trace_id = %s, span trace_id = %s', self._trace.id, span.trace_id) 137 return 138 if not self._trace.stack: 139 log('warning: finish span called but stack is empty') 140 return 141 if self._trace.stack[-1].id != span.id: 142 log('warning: finished span is not the currently active span') 143 return 144 self._trace.stack.pop() 145 146 def finish_trace(self, span): 147 self.finish_span(span) 148 self._trace = None 149 150 def parse_http_trace(self, request): 151 if not self.http_trace_parser_hook: 152 return None 153 return self.http_trace_parser_hook(request) 154 155 def propagate_and_start_trace(self, context, request): 156 err = None 157 propagation_context = None 158 try: 159 propagation_context = self.parse_http_trace(request) 160 except Exception: 161 err = sys.exc_info()[0] 162 log('error: http_trace_parser_hook returned exception: %s', sys.exc_info()[0]) 163 164 if propagation_context: 165 return self.start_trace(context=context, trace_id=propagation_context.trace_id, parent_span_id=propagation_context.parent_id, dataset=propagation_context.dataset) 166 for k, v in propagation_context.trace_fields: 167 self.add_trace_field(k, v) 168 169 else: 170 if err is not None: 171 context['parser_hook_error'] = repr(err) 172 return self.start_trace(context, trace_id=None, parent_span_id=None) 173 174 def get_propagation_context(self): 175 if not self._trace: 176 return None 177 return beeline.propagation.PropagationContext(self.get_active_trace_id(), self.get_active_span().id, self._trace.fields) 178 179 def get_active_trace_id(self): 180 if self._trace: 181 return self._trace.id 182 183 def get_active_span(self): 184 if self._trace and self._trace.stack: 185 return self._trace.stack[-1] 186 187 def add_context_field(self, name, value): 188 span = self.get_active_span() 189 if span: 190 span.add_context_field(name=name, value=value) 191 192 def add_context(self, data): 193 span = self.get_active_span() 194 if span: 195 span.add_context(data=data) 196 197 def remove_context_field(self, name): 198 span = self.get_active_span() 199 if span: 200 span.remove_context_field(name=name) 201 202 def add_rollup_field(self, name, value): 203 value = float(value) 204 span = self.get_active_span() 205 if span: 206 span.rollup_fields[name] += value 207 if not self._trace: 208 log('warning: adding rollup field without an active trace') 209 return 210 self._trace.rollup_fields['rollup.%s' % name] += value 211 212 def add_trace_field(self, name, value): 213 if type(name) == str and not name.startswith('app.') or type(name) != str: 214 key = 'app.%s' % name 215 else: 216 key = name 217 self.add_context_field(key, value) 218 if not self._trace: 219 log('warning: adding trace field without an active trace') 220 return 221 self._trace.fields[key] = value 222 223 def remove_trace_field(self, name): 224 key = 'app.%s' % name 225 self.remove_context_field(key) 226 if not self._trace: 227 log('warning: removing trace field without an active trace') 228 return 229 self._trace.fields.pop(key) 230 231 def marshal_trace_context(self): 232 if not self._trace: 233 log('warning: marshal_trace_context called, but no active trace') 234 return 235 return marshal_trace_context(self._trace.id, self._trace.stack[-1].id, self._trace.fields) 236 237 def register_hooks(self, presend = None, sampler = None, http_trace_parser = None, http_trace_propagation = None): 238 self.presend_hook = presend 239 self.sampler_hook = sampler 240 self.http_trace_parser_hook = http_trace_parser 241 self.http_trace_propagation_hook = http_trace_propagation 242 243 def _run_hooks_and_send(self, span): 244 presampled = False 245 if self.sampler_hook: 246 log('executing sampler hook on event ev = %s', span.event.fields()) 247 keep, new_rate = self.sampler_hook(span.event.fields()) 248 if not keep: 249 log('skipping event due to sampler hook sampling ev = %s', span.event.fields()) 250 return 251 span.event.sample_rate = new_rate 252 presampled = True 253 if self.presend_hook: 254 log('executing presend hook on event ev = %s', span.event.fields()) 255 self.presend_hook(span.event.fields()) 256 if presampled: 257 log('enqueuing presampled event ev = %s', span.event.fields()) 258 span.event.send_presampled() 259 elif _should_sample(span.trace_id, span.event.sample_rate): 260 span.event.send_presampled() 261 262 263 class SynchronousTracer(Tracer): 264 265 def __init__(self, client): 266 super(SynchronousTracer, self).__init__(client) 267 self._state = threading.local() 268 269 @property 270 def _trace(self): 271 return getattr(self._state, 'trace', None) 272 273 @_trace.setter 274 def _trace(self, new_trace): 275 self._state.trace = new_trace 276 277 278 class Span(object): 279 280 def __init__(self, trace_id, parent_id, id, event, is_root = False): 281 self.trace_id = trace_id 282 self.parent_id = parent_id 283 self.id = id 284 self.event = event 285 self.event.start_time = datetime.datetime.now() 286 self.rollup_fields = defaultdict(float) 287 self._is_root = is_root 288 289 def add_context_field(self, name, value): 290 self.event.add_field(name, value) 291 292 def add_context(self, data): 293 self.event.add(data) 294 295 def remove_context_field(self, name): 296 if name in self.event.fields(): 297 del self.event.fields()[name] 298 299 def is_root(self): 300 return self._is_root 301 302 303 def _should_sample(trace_id, sample_rate): 304 sample_upper_bound = MAX_INT32 / sample_rate 305 sha1 = hashlib.sha1() 306 sha1.update(trace_id.encode('utf-8')) 307 value, = struct.unpack('>I', sha1.digest()[:4]) 308 if value < sample_upper_bound: 309 return True 310 return False 311 312 313 def marshal_trace_context(trace_id, parent_id, context): 314 version = 1 315 trace_fields = base64.b64encode(json.dumps(context).encode()).decode() 316 trace_context = '{};trace_id={},parent_id={},context={}'.format(version, trace_id, parent_id, trace_fields) 317 return trace_context 318 319 320 def unmarshal_trace_context(trace_header): 321 return beeline.propagation.honeycomb.unmarshal_propagation_context(trace_header) 322 323 324 def traced_impl(tracer_fn, name, trace_id, parent_id): 325 326 def wrapped(fn): 327 if inspect.isgeneratorfunction(fn): 328 329 @functools.wraps(fn) 330 def inner(*args, **kwargs): 331 inner_generator = fn(*args, **kwargs) 332 with tracer_fn(name=name, trace_id=trace_id, parent_id=parent_id): 333 for value in inner_generator: 334 yield value 335 336 return inner 337 else: 338 339 @functools.wraps(fn) 340 def inner(*args, **kwargs): 341 with tracer_fn(name=name, trace_id=trace_id, parent_id=parent_id): 342 return fn(*args, **kwargs) 343 344 return inner 345 346 return wrapped 347 348 349 system_random = random.SystemRandom() 350 351 def generate_span_id(): 352 format_str = '{{:0{:d}x}}'.format(SPAN_ID_BYTES * 2) 353 return format_str.format(system_random.getrandbits(SPAN_ID_BYTES * 8)) 354 355 356 def generate_trace_id(): 357 format_str = '{{:0{:d}x}}'.format(TRACE_ID_BYTES * 2) 358 return format_str.format(system_random.getrandbits(TRACE_ID_BYTES * 8))