handle health checks in middleware

Author: andrewsg

multicore_runtime/middleware.py

@@ -17,6 +17,9 @@
 import logging
 import os
 
+from werkzeug.wrappers import Request
+from werkzeug.wrappers import Response
+
 
 RESERVED_ENV_KEYS = {
     'AUTH_DOMAIN': 'gmail.com',  # Default auth domain must be set.
@@ -160,3 +163,23 @@ def get_env_to_hide_service_bridge(wsgi_env):
                     'SERVER_PORT', https)
 
   return output
+
+
+def health_check_middleware(app):
+  """Intercept requests to /_ah/health and respond healthy (200 OK).
+
+  Args:
+    app: The WSGI app to wrap.
+
+  Returns:
+    The wrapped app, also a WSGI app.
+  """
+
+  def health_check_intercept_wrapper(wsgi_env, start_response):  # pylint: disable=missing-docstring
+    request = Request(wsgi_env)
+    if request.path == '/_ah/health':  # Only intercept exact matches.
+      return Response('healthy', status=200)(wsgi_env, start_response)
+    else:
+      return app(wsgi_env, start_response)
+
+  return health_check_intercept_wrapper

multicore_runtime/wsgi.py

@@ -25,6 +25,7 @@
 
 from dispatcher import dispatcher
 from middleware import reset_environment_middleware
+from middleware import health_check_middleware
 from wsgi_config import env_vars_from_env_config
 from wsgi_config import get_module_config
 from wsgi_config import get_module_config_filename
@@ -83,7 +84,11 @@
 # layer (the middleware code that will process a request first). Inside the
 # innermost layer is the actual dispatcher, above.
 
+# Intercept health check requests on /_ah/health. This is a temporary measure
+# until container-level health check handlers are in place and turned on.
+meta_app = health_check_middleware(meta_app)
+
 # Reset os.environ to the frozen state and add request-specific data.
 meta_app = reset_environment_middleware(meta_app, frozen_environment,
                                         frozen_user_env,
-                                        frozen_env_config_env)
+                                        frozen_env_config_env)