From a2edb38a8595ec7d1551e951b43d95836a0e0ff8 Mon Sep 17 00:00:00 2001
From: Mike Lang <mikelang3000@gmail.com>
Date: Sat, 29 Jun 2019 18:36:13 -0700
Subject: [PATCH] Add an InstrumentedSession wrapper that automatically gathers
 metrics on http calls

---
 common/common/requests.py | 51 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)
 create mode 100644 common/common/requests.py

diff --git a/common/common/requests.py b/common/common/requests.py
new file mode 100644
index 0000000..59896e5
--- /dev/null
+++ b/common/common/requests.py
@@ -0,0 +1,51 @@
+
+"""Code for instrumenting requests calls. Requires requests, obviously."""
+
+import urlparse
+
+import requests
+import prometheus_client as prom
+from monotonic import monotonic
+
+request_latency = prom.Histogram(
+	'http_client_request_latency',
+	'Time taken to make an outgoing HTTP request. '
+	'Status = "error" is used if an error occurs. Measured as time from first byte sent to '
+	'headers finished being parsed, ie. does not include reading a streaming response.',
+	['name', 'method', 'domain', 'status'],
+)
+
+response_size = prom.Histogram(
+	'http_client_response_size',
+	"The content length of (non-streaming) responses to outgoing HTTP requests.",
+	['name', 'method', 'domain', 'status'],
+)
+
+request_concurrency = prom.Gauge(
+	'http_client_request_concurrency',
+	"The number of outgoing HTTP requests currently ongoing",
+	['name', 'method', 'domain'],
+)
+
+class InstrumentedSession(requests.Session):
+	"""A requests Session that automatically records metrics on requests made.
+	Users may optionally pass a 'metric_name' kwarg that will be included as the 'name' label.
+	"""
+
+	def request(self, method, url, *args, **kwargs):
+		_, domain, _, _, _ = urlparse.urlsplit(url)
+		name = kwargs.pop('metric_name', '')
+
+		start = monotonic() # we only use our own measured latency if an error occurs
+		try:
+			with request_concurrency.labels(name, method, domain).track_inprogress():
+				response = super(InstrumentedSession, self).request(method, url, *args, **kwargs)
+		except Exception:
+			latency = monotonic() - start
+			request_latency.labels(name, method, domain, "error").observe(latency)
+			raise
+
+		request_latency.labels(name, method, domain, response.status_code).observe(response.elapsed)
+		if 'content-length' in response.headers:
+			response_size.labels(name, method, domain, response.status_code).observe(response.headers['content-length'])
+		return response