Health check responder (#2639)

* Revert "Revert "Health check responder (#2636)" (#2638)"

This reverts commit ae4867a.

* Revert "Include the health check reponser into heartbeat"

This reverts commit 3834b08.

* Include the health check reponser into run_heartbeat

* Allow outside network connections to get in.

* Run responder server in its own thread

* fix http import

* map docker port to host port

* import with full path

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

local/run_docker.bash

@@ -46,4 +46,5 @@ sudo docker run -e COMMAND_OVERRIDE="$COMMAND_OVERRIDE" -e SETUP_NFS= -e HOST_UI
               -e LOCAL_SRC=$LOCAL_SRC \
               -e CONFIG_DIR_OVERRIDE=$CONFIG_DIR_OVERRIDE \
               --hostname test-bot-$USER \
+              -p 7123:7123 \
               -ti --privileged --cap-add=all $IMAGE "$@"

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,52 @@
+# 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."""
+
+from http import HTTPStatus
+from http.server import BaseHTTPRequestHandler
+from http.server import HTTPServer
+import threading
+
+from clusterfuzz._internal.system import process_handler
+
+RESPONDER_IP = '0.0.0.0'
+RESPONDER_PORT = 7123
+EXPECTED_SCRIPTS = ['run.py', 'run_bot.py']
+
+
+class RequestHandler(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 = HTTPStatus.OK
+    else:
+      response_code = 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 = HTTPServer((RESPONDER_IP, RESPONDER_PORT),
+                                             RequestHandler)
+  server_thread = threading.Thread(
+      target=health_check_responder_server.serve_forever)
+  server_thread.start()

src/python/bot/startup/run_heartbeat.py

@@ -29,6 +29,8 @@
 from clusterfuzz._internal.metrics import logs
 from clusterfuzz._internal.system import environment
 from clusterfuzz._internal.system import shell
+from python.bot.startup.health_check_responder import \
+    run_server as run_health_responser_server
 
 BEAT_SCRIPT = 'heartbeat.py'
 
@@ -51,6 +53,8 @@ def main():
   beat_interpreter = shell.get_interpreter(beat_script_path)
   assert beat_interpreter
 
+  run_health_responser_server()
+
   while True:
     beat_command = [
         beat_interpreter, beat_script_path,