added initial telemetry work

rollbar/__init__.py

@@ -20,7 +20,7 @@
 import requests
 import six
 
-from rollbar.lib import events, filters, dict_merge, parse_qs, text, transport, urljoin, iteritems, defaultJSONEncode
+from rollbar.lib import events, filters, dict_merge, parse_qs, text, transport, telemetry, urljoin, iteritems, defaultJSONEncode
 
 
 __version__ = '0.16.1'
@@ -381,6 +381,22 @@ def init(access_token, environment='production', scrub_fields=None, url_fields=N
     if SETTINGS.get('allow_logging_basic_config'):
         logging.basicConfig()
 
+    queue_size = SETTINGS.get('set_custom_queue_size')
+    if queue_size:
+        telemetry.TELEMETRY_QUEUE = telemetry.Queue(queue_size)
+
+    if SETTINGS.get('log_telemetry'):
+        formatter = SETTINGS.get('log_telemetry_formatter')
+        telemetry.set_log_telemetry(formatter)
+    if SETTINGS.get('network_telemetry'):
+        enable_req_headers = SETTINGS.get('enable_req_headers')
+        enable_response_headers = SETTINGS.get('enable_response_headers')
+        requests.get = telemetry.request(requests.get, enable_req_headers, enable_response_headers )
+        requests.post = telemetry.request(requests.post, enable_req_headers, enable_response_headers)
+        requests.put = telemetry.request(requests.put, enable_req_headers, enable_response_headers)
+        requests.patch = telemetry.request(requests.patch, enable_req_headers, enable_response_headers)
+        requests.delete = telemetry.request(requests.delete, enable_req_headers, enable_response_headers)
+
     if SETTINGS.get('handler') == 'agent':
         agent_log = _create_agent_log()
 
@@ -777,6 +793,7 @@ def _report_exc_info(exc_info, request, extra_data, payload_data, level=None):
     _add_request_data(data, request)
     _add_person_data(data, request)
     _add_lambda_context_data(data)
+    _add_telemetry(data)
     data['server'] = _build_server_data()
 
     if payload_data:
@@ -857,6 +874,7 @@ def _report_message(message, level, request, extra_data, payload_data):
     _add_request_data(data, request)
     _add_person_data(data, request)
     _add_lambda_context_data(data)
+    _add_telemetry(data)
     data['server'] = _build_server_data()
 
     if payload_data:
@@ -1119,6 +1137,12 @@ def _add_request_data(data, request):
             data['request'] = request_data
 
 
+def _add_telemetry(data):
+    telemetry_data = telemetry.TELEMETRY_QUEUE.get_items()
+    if telemetry_data:
+        data['body']['telemetry'] = telemetry_data
+
+
 def _check_add_locals(frame, frame_num, total_frames):
     """
     Returns True if we should record local variables for the given frame.

rollbar/lib/telemetry.py

@@ -0,0 +1,82 @@
+import logging
+import threading
+import time
+
+import rollbar
+
+
+class Queue:
+    def __init__(self, size):
+        self.size = size
+        self.items = []
+        self.lock = threading.Lock()
+
+    def put(self, item):
+        with self.lock:
+            if len(self.items) >= self.size:
+                self.items = self.items[1:]
+            self.items.append(item)
+
+    def get_items(self):
+        with self.lock:
+            return self.items
+
+
+TELEMETRY_QUEUE_SIZE = 50
+TELEMETRY_QUEUE = Queue(TELEMETRY_QUEUE_SIZE)
+
+
+class TelemetryLogHandler(logging.Handler):
+    def __init__(self, formatter=None):
+        super(TelemetryLogHandler, self).__init__()
+        self.formatter = formatter
+
+    def emit(self, record):
+        self.setFormatter(self.formatter)
+        msg = {'message': self.format(record)}
+        data = {'body': msg, 'source': 'client', 'timestamp_ms': int(time.time()), 'type': 'log',
+                'level': record.levelname}
+
+        TELEMETRY_QUEUE.put(data)
+
+
+def set_log_telemetry(log_formatter=None):
+    logging.getLogger().addHandler(TelemetryLogHandler(log_formatter))
+
+
+def request(request_function, enable_req_headers, enable_response_headers):
+    def telemetry(*args, **kwargs):
+
+        def clean_headers(headers):
+            for h in headers:
+                if h in rollbar.SETTINGS['scrub_fields']:
+                    del headers[h]
+            return headers
+
+        data = {'level': 'info'}
+        data_body = {'status_code': None}
+
+        response = request_function(*args, **kwargs)
+
+        if response:
+            data_body['status_code'] = response.status_code
+            if response.status_code >= 500:
+                data['level'] = 'critical'
+            elif response.status_code >= 400:
+                data['level'] = 'error'
+            if enable_response_headers:
+                data_body['response'] = {'headers': clean_headers(response.headers)}
+        if enable_req_headers:
+            data_body['request_headers'] = clean_headers(kwargs.get('headers'))
+        data_body['url'] = args[0]
+        data_body['method'] = request_function.__name__
+        data_body['subtype'] = 'http'
+        data['body'] = data_body
+
+        data['source'] = 'client'
+        data['timestamp_ms'] = int(time.time())
+        data['type'] = 'network'
+        TELEMETRY_QUEUE.put(data)
+
+        return response
+    return telemetry