Health check responder (#2636)

* A responder to health checks

* Include the health check responder in `heartbeat.py`

* A note about checking `run_bot.py`

* Testcases for health check responder

* Allow health check packets in configs

Author: DonggeLiu

configs/test/gce/linux-init-ml-with-gpu.yaml

@@ -31,6 +31,7 @@ bootcmd:
   - echo never > /sys/kernel/mm/transparent_hugepage/defrag
   - echo core > /proc/sys/kernel/core_pattern # for AFL
   - swapon -a
+  - iptables -w -A INPUT -p tcp --dport 7123 -j ACCEPT # health check port
 
   # Note that NVIDIA_DRIVER_VERSION can be used with a particular CUDA version
   # only, e.g. CUDA 9.0 needs drivers version 384.130. CUDA and TensorFlow are

configs/test/gce/linux-init.yaml

@@ -31,6 +31,7 @@ bootcmd:
   - echo never > /sys/kernel/mm/transparent_hugepage/defrag
   - echo core > /proc/sys/kernel/core_pattern # for AFL
   - swapon -a
+  - iptables -w -A INPUT -p tcp --dport 7123 -j ACCEPT # health check port
 
 write_files:
   - path: /etc/systemd/system/clusterfuzz.service

src/clusterfuzz/_internal/system/process_handler.py

@@ -649,3 +649,16 @@ def terminate_processes_matching_cmd_line(match_strings,
     if any(x in process_path for x in match_strings):
       if not any([x in process_path for x in exclude_strings]):
         terminate_process(process_info['pid'], kill)
+
+
+def scripts_are_running(expected_scripts):
+  """Check if all target scripts are running as expected."""
+  scripts_left = expected_scripts.copy()
+  for process in psutil.process_iter():
+    for expected_script in scripts_left:
+      if any(expected_script == os.path.basename(cmdline)
+             for cmdline in process.cmdline()):
+        scripts_left.remove(expected_script)
+        if not scripts_left:
+          return True
+  return False

src/clusterfuzz/_internal/tests/core/bot/startup/health_check_responder_test.py

@@ -0,0 +1,84 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""health check reposnser tests."""
+
+from http.server import HTTPServer
+import threading
+import unittest
+
+import mock
+import requests
+
+from python.bot.startup.health_check_responder import EXPECTED_SCRIPTS
+from python.bot.startup.health_check_responder import RequestHandler
+from python.bot.startup.health_check_responder import RESPONDER_IP
+from python.bot.startup.health_check_responder import RESPONDER_PORT
+
+RESPONDER_ADDR = f'http://{RESPONDER_IP}:{RESPONDER_PORT}'
+
+
+class HealthCheckResponderTest(unittest.TestCase):
+  """Test health check responder."""
+
+  def setUp(self):
+    """Prepare mock processes and start the responder server thread."""
+    self.mock_run_process = mock.MagicMock()
+    self.mock_run_process.cmdline.return_value = ['./' + EXPECTED_SCRIPTS[0]]
+    self.mock_run_bot_process = mock.MagicMock()
+    self.mock_run_bot_process.cmdline.return_value = [
+        './' + EXPECTED_SCRIPTS[1]
+    ]
+
+    self.health_check_responder_server = HTTPServer(
+        (RESPONDER_IP, RESPONDER_PORT), RequestHandler)
+    server_thread = threading.Thread(
+        target=self.health_check_responder_server.serve_forever)
+    server_thread.start()
+
+  def tearDown(self):
+    self.health_check_responder_server.shutdown()
+    self.health_check_responder_server.server_close()
+
+  @mock.patch(
+      'python.bot.startup.health_check_responder.process_handler.psutil')
+  def test_healthy(self, mock_psutil):
+    """Testcase for both scripts are running."""
+    mock_psutil.process_iter.return_value = [
+        self.mock_run_process, self.mock_run_bot_process
+    ]
+
+    self.assertEqual(200, requests.get(f'{RESPONDER_ADDR}').status_code)
+
+  @mock.patch(
+      'python.bot.startup.health_check_responder.process_handler.psutil')
+  def test_run_terminated(self, mock_psutil):
+    """Testcase for only the run script is running."""
+    mock_psutil.process_iter.return_value = [self.mock_run_process]
+
+    self.assertEqual(500, requests.get(f'{RESPONDER_ADDR}').status_code)
+
+  @mock.patch(
+      'python.bot.startup.health_check_responder.process_handler.psutil')
+  def test_run_bot_terminated(self, mock_psutil):
+    """Testcase for only the run_bot script is running."""
+    mock_psutil.process_iter.return_value = [self.mock_run_bot_process]
+
+    self.assertEqual(500, requests.get(f'{RESPONDER_ADDR}').status_code)
+
+  @mock.patch(
+      'python.bot.startup.health_check_responder.process_handler.psutil')
+  def test_both_terminated(self, mock_psutil):
+    """Testcase for neither script is running."""
+    mock_psutil.process_iter.return_value = []
+    self.assertEqual(500, requests.get(f'{RESPONDER_ADDR}').status_code)

src/python/bot/startup/health_check_responder.py

@@ -0,0 +1,46 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Health check responder that checks if all scripts are running as expected
+   and responds to health checks."""
+import http
+
+from clusterfuzz._internal.system import process_handler
+
+RESPONDER_IP = 'localhost'
+RESPONDER_PORT = 7123
+EXPECTED_SCRIPTS = ['run.py', 'run_bot.py']
+
+
+class RequestHandler(http.server.BaseHTTPRequestHandler):
+  """Handler for GET request form the health checker."""
+
+  def do_GET(self):  # pylint: disable=invalid-name
+    """Handle a GET request."""
+    if process_handler.scripts_are_running(EXPECTED_SCRIPTS):
+      # Note: run_bot.py is expected to go down during source updates
+      #   (which can take a few minutes)
+      # Health checks should be resilient to this
+      # and set a threshold / check interval to account for this.
+      response_code = http.HTTPStatus.OK
+    else:
+      response_code = http.HTTPStatus.INTERNAL_SERVER_ERROR
+    self.send_response(response_code)
+    self.end_headers()
+
+
+def run_server():
+  """Start a HTTP server to respond to the health checker."""
+  health_check_responder_server = http.server.HTTPServer(
+      (RESPONDER_IP, RESPONDER_PORT), RequestHandler)
+  health_check_responder_server.serve_forever()

src/python/bot/startup/heartbeat.py

@@ -25,6 +25,8 @@
 import sys
 import time
 
+from health_check_responder import run_server as run_health_responder_server
+
 from clusterfuzz._internal.base import dates
 from clusterfuzz._internal.base import tasks
 from clusterfuzz._internal.datastore import data_handler
@@ -113,6 +115,8 @@ def main():
   except Exception:
     logs.log_error('Failed to beat.')
 
+  run_health_responder_server()
+
   time.sleep(data_types.HEARTBEAT_WAIT_INTERVAL)