From b0ded641c34c81de73703cb24e138a84327878b5 Mon Sep 17 00:00:00 2001
From: Mike Lang <mikelang3000@gmail.com>
Date: Sun, 6 Jan 2019 00:14:39 -0800
Subject: [PATCH] Add a logging handler which counts logs for prometheus stats

This isn't as good as having a full centralised logging system, but should
suffice to know if anything funny is happening.
---
 backfiller/backfiller/main.py |  2 ++
 common/common.py              | 14 ++++++++++++++
 common/setup.py               |  2 +-
 downloader/downloader/main.py |  1 +
 restreamer/restreamer/main.py |  4 +++-
 5 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/backfiller/backfiller/main.py b/backfiller/backfiller/main.py
index 666b9cf..83bda5c 100644
--- a/backfiller/backfiller/main.py
+++ b/backfiller/backfiller/main.py
@@ -234,6 +234,8 @@ def main(base_dir='.', stream='', variants='', fill_wait=5, full_fill_wait=180,
 
 	variants = variants.split(',') if variants else []
 
+	common.PromLogCountsHandler.install()
+
 	logging.info('Starting backfilling {} with {} as variants to {}'.format(stream, ', '.join(variants), base_dir))
 
 	fill_start = datetime.datetime.now()
diff --git a/common/common.py b/common/common.py
index 356bee4..451370e 100644
--- a/common/common.py
+++ b/common/common.py
@@ -13,6 +13,7 @@ import sys
 from collections import namedtuple
 
 import dateutil.parser
+import prometheus_client as prom
 
 
 def dt_to_bustime(start, dt):
@@ -299,3 +300,16 @@ def encode_strings(o):
 	if isinstance(o, unicode):
 		return o.encode('utf-8')
 	return o
+
+
+log_count = prom.Counter("log_count", "Count of messages logged", ["level", "module", "function"])
+
+class PromLogCountsHandler(logging.Handler):
+	"""A logging handler that records a count of logs by level, module and function."""
+	def emit(self, record):
+		log_count.labels(record.levelname, record.module, record.funcName).inc()
+
+	@classmethod
+	def install(cls):
+		root_logger = logging.getLogger()
+		root_logger.addHandler(cls())
diff --git a/common/setup.py b/common/setup.py
index 90bc79f..05618a7 100644
--- a/common/setup.py
+++ b/common/setup.py
@@ -5,7 +5,7 @@ setup(
 	version = "0.0.0",
 	py_modules = ["common"],
 	install_requires = [
+		"prometheus-client",
 		"python-dateutil",
-		"PyYAML<4.0.0",
 	],
 )
diff --git a/downloader/downloader/main.py b/downloader/downloader/main.py
index 0aac4d2..455c60e 100644
--- a/downloader/downloader/main.py
+++ b/downloader/downloader/main.py
@@ -494,6 +494,7 @@ def main(channel, base_dir=".", qualities=""):
 	qualities = qualities.split(",") if qualities else []
 	manager = StreamsManager(channel, base_dir, qualities)
 	gevent.signal(signal.SIGTERM, manager.stop) # shut down on sigterm
+	common.PromLogCountsHandler.install()
 	logging.info("Starting up")
 	manager.run()
 	logging.info("Gracefully stopped")
diff --git a/restreamer/restreamer/main.py b/restreamer/restreamer/main.py
index 429cd9a..c08a902 100644
--- a/restreamer/restreamer/main.py
+++ b/restreamer/restreamer/main.py
@@ -16,7 +16,7 @@ from flask import Flask, url_for, request, abort, Response
 from gevent import subprocess
 from gevent.pywsgi import WSGIServer
 
-from common import get_best_segments
+from common import get_best_segments, PromLogCountsHandler
 
 import generate_hls
 from stats import stats, after_request
@@ -397,6 +397,8 @@ def main(host='0.0.0.0', port=8000, base_dir='.'):
 		server.stop()
 	gevent.signal(signal.SIGTERM, stop)
 
+	PromLogCountsHandler.install()
+
 	logging.info("Starting up")
 	server.serve_forever()
 	logging.info("Gracefully shut down")