base.py (6730B)
1 #Embedded file name: C:\BuildAgent\work\20d73f1b24eed06e\eve\release\V24.01\packages\datadog\dogstatsd\base.py 2 from random import random 3 import logging 4 import os 5 import socket 6 from datadog.dogstatsd.context import TimedContextManagerDecorator 7 from datadog.dogstatsd.route import get_default_route 8 from datadog.util.compat import imap, text 9 log = logging.getLogger('datadog.dogstatsd') 10 11 class DogStatsd(object): 12 OK, WARNING, CRITICAL, UNKNOWN = (0, 1, 2, 3) 13 14 def __init__(self, host = 'localhost', port = 8125, max_buffer_size = 50, namespace = None, constant_tags = None, use_ms = False, use_default_route = False): 15 self.host = self.resolve_host(host, use_default_route) 16 self.port = int(port) 17 self.socket = None 18 self.max_buffer_size = max_buffer_size 19 self._send = self._send_to_server 20 self.encoding = 'utf-8' 21 env_tags = [ tag for tag in os.environ.get('DATADOG_TAGS', '').split(',') if tag ] 22 if constant_tags is None: 23 constant_tags = [] 24 self.constant_tags = constant_tags + env_tags 25 self.namespace = namespace 26 self.use_ms = use_ms 27 28 def __enter__(self): 29 self.open_buffer(self.max_buffer_size) 30 return self 31 32 def __exit__(self, type, value, traceback): 33 self.close_buffer() 34 35 @staticmethod 36 def resolve_host(host, use_default_route): 37 if not use_default_route: 38 return host 39 return get_default_route() 40 41 def get_socket(self): 42 if not self.socket: 43 sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) 44 sock.connect((self.host, self.port)) 45 self.socket = sock 46 return self.socket 47 48 def open_buffer(self, max_buffer_size = 50): 49 self.max_buffer_size = max_buffer_size 50 self.buffer = [] 51 self._send = self._send_to_buffer 52 53 def close_buffer(self): 54 self._send = self._send_to_server 55 self._flush_buffer() 56 57 def gauge(self, metric, value, tags = None, sample_rate = 1): 58 return self._report(metric, 'g', value, tags, sample_rate) 59 60 def increment(self, metric, value = 1, tags = None, sample_rate = 1): 61 self._report(metric, 'c', value, tags, sample_rate) 62 63 def decrement(self, metric, value = 1, tags = None, sample_rate = 1): 64 metric_value = -value if value else value 65 self._report(metric, 'c', metric_value, tags, sample_rate) 66 67 def histogram(self, metric, value, tags = None, sample_rate = 1): 68 self._report(metric, 'h', value, tags, sample_rate) 69 70 def timing(self, metric, value, tags = None, sample_rate = 1): 71 self._report(metric, 'ms', value, tags, sample_rate) 72 73 def timed(self, metric = None, tags = None, sample_rate = 1, use_ms = None): 74 return TimedContextManagerDecorator(self, metric, tags, sample_rate, use_ms) 75 76 def set(self, metric, value, tags = None, sample_rate = 1): 77 self._report(metric, 's', value, tags, sample_rate) 78 79 def close_socket(self): 80 if self.socket: 81 self.socket.close() 82 self.socket = None 83 84 def _report(self, metric, metric_type, value, tags, sample_rate): 85 if value is None: 86 return 87 if sample_rate != 1 and random() > sample_rate: 88 return 89 payload = [] 90 if self.constant_tags: 91 if tags: 92 tags = tags + self.constant_tags 93 else: 94 tags = self.constant_tags 95 if self.namespace: 96 payload.extend([self.namespace, '.']) 97 payload.extend([metric, 98 ':', 99 value, 100 '|', 101 metric_type]) 102 if sample_rate != 1: 103 payload.extend(['|@', sample_rate]) 104 if tags: 105 payload.extend(['|#', ','.join(tags)]) 106 encoded = ''.join(imap(text, payload)) 107 self._send(encoded) 108 109 def _send_to_server(self, packet): 110 try: 111 (self.socket or self.get_socket()).send(packet.encode(self.encoding)) 112 except socket.error: 113 log.info('Error submitting packet, will try refreshing the socket') 114 self.close_socket() 115 try: 116 self.get_socket().send(packet.encode(self.encoding)) 117 except socket.error: 118 self.close_socket() 119 log.exception('Failed to send packet with a newly bound socket') 120 121 def _send_to_buffer(self, packet): 122 self.buffer.append(packet) 123 if len(self.buffer) >= self.max_buffer_size: 124 self._flush_buffer() 125 126 def _flush_buffer(self): 127 self._send_to_server('\n'.join(self.buffer)) 128 self.buffer = [] 129 130 def _escape_event_content(self, string): 131 return string.replace('\n', '\\n') 132 133 def _escape_service_check_message(self, string): 134 return string.replace('\n', '\\n').replace('m:', 'm\\:') 135 136 def event(self, title, text, alert_type = None, aggregation_key = None, source_type_name = None, date_happened = None, priority = None, tags = None, hostname = None): 137 title = self._escape_event_content(title) 138 text = self._escape_event_content(text) 139 if self.constant_tags: 140 if tags: 141 tags += self.constant_tags 142 else: 143 tags = self.constant_tags 144 string = u'_e{%d,%d}:%s|%s' % (len(title), 145 len(text), 146 title, 147 text) 148 if date_happened: 149 string = '%s|d:%d' % (string, date_happened) 150 if hostname: 151 string = '%s|h:%s' % (string, hostname) 152 if aggregation_key: 153 string = '%s|k:%s' % (string, aggregation_key) 154 if priority: 155 string = '%s|p:%s' % (string, priority) 156 if source_type_name: 157 string = '%s|s:%s' % (string, source_type_name) 158 if alert_type: 159 string = '%s|t:%s' % (string, alert_type) 160 if tags: 161 string = '%s|#%s' % (string, ','.join(tags)) 162 if len(string) > 8192: 163 raise Exception(u'Event "%s" payload is too big (more that 8KB), event discarded' % title) 164 self._send(string) 165 166 def service_check(self, check_name, status, tags = None, timestamp = None, hostname = None, message = None): 167 message = self._escape_service_check_message(message) if message is not None else '' 168 string = u'_sc|{0}|{1}'.format(check_name, status) 169 if timestamp: 170 string = u'{0}|d:{1}'.format(string, timestamp) 171 if hostname: 172 string = u'{0}|h:{1}'.format(string, hostname) 173 if tags: 174 string = u'{0}|#{1}'.format(string, ','.join(tags)) 175 if message: 176 string = u'{0}|m:{1}'.format(string, message) 177 self._send(string) 178 179 180 statsd = DogStatsd()