From e9f129ac8478d5f8b95ebc803186f2526e5639be Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniele=20Vigan=C3=B2?= <daniele@vigano.me>
Date: Wed, 5 Jan 2022 12:19:42 +0100
Subject: [PATCH] Enable HC on driver

---
 dcc/driver/apps.py   |  5 +++++
 dcc/driver/health.py | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+)
 create mode 100644 dcc/driver/health.py

diff --git a/dcc/driver/apps.py b/dcc/driver/apps.py
index 4a5df03..27eca5a 100644
--- a/dcc/driver/apps.py
+++ b/dcc/driver/apps.py
@@ -1,6 +1,11 @@
 from django.apps import AppConfig
+from health_check.plugins import plugin_dir
 
 
 class DriverConfig(AppConfig):
     default_auto_field = 'django.db.models.BigAutoField'
     name = 'driver'
+
+    def ready(self):
+        from driver.health import DriverHealthCheck
+        plugin_dir.register(DriverHealthCheck)
diff --git a/dcc/driver/health.py b/dcc/driver/health.py
new file mode 100644
index 0000000..75e09f9
--- /dev/null
+++ b/dcc/driver/health.py
@@ -0,0 +1,20 @@
+from health_check.backends import BaseHealthCheckBackend
+from health_check.exceptions import (ServiceUnavailable,
+                                     ServiceReturnedUnexpectedResult)
+
+from driver.connector import Connector
+
+
+class DriverHealthCheck(BaseHealthCheckBackend):
+    critical_service = False
+
+    def check_status(self):
+        try:
+            Connector().passthrough(b'<s>')
+        except ConnectionRefusedError as e:
+            self.add_error(ServiceUnavailable("IOError"), e)
+        except Exception as e:
+            self.add_error(ServiceReturnedUnexpectedResult("IOError"), e)
+
+    def identifier(self):
+        return "DriverDaemon"