#28 - implement a safe shutdown endpoint

Author: dev-dull

constfig.py

@@ -20,6 +20,10 @@ def __init__(self):
         self.load_config()
 
     def set_constants(self):
+        # `C` (the instance of `_C()` set below) is only created once even though it is imported in `ddb` and `pyxie`.
+        # Probably not a good idea for _SHUTDOWN to depend on that behavior for the long term. Only used in `pyxie` currently.
+        self._SHUTDOWN = False  # Toggled to true to cleanly stop the service and ensure data is persisted to disk.
+
         # Constant values
         self.APP_NAME = "pyxie"
         self.LOG = logging.getLogger(self.APP_NAME)

pyxie.py

@@ -1,7 +1,8 @@
+import os
 from ddb import DDB
-from time import time
+from time import time, sleep
 from constfig import C
-from flask import Flask, Response, request
+from flask import Flask, Response, request, abort
 
 
 pyxie = Flask(C.APP_NAME)
@@ -13,39 +14,54 @@ def _validate_api_key():
     return api_key in C.API_KEYS
 
 
+@pyxie.before_request
+def check_if_shutting_down():
+    # Since we're currently single threaded, we shouldn't really need this check. This is here
+    # for any weird edge cases with flask.
+    if C._SHUTDOWN:
+        abort(503, description="Service is shutting down. Please try again later.")
+
+
+@pyxie.before_request
+def validate_api_key():
+    if request.path == "/" or _validate_api_key():
+        return
+    return "Unauthorized", 401
+
+
+@pyxie.after_request
+def shutdown(response):
+    if C._SHUTDOWN:
+        _data.dump()
+        os._exit(0)
+    return response
+
+
 @pyxie.route("/register", methods=[C.HTTP_METHOD_POST])
 def register():
-    if _validate_api_key():
-        _data.register()
-        return "Success", 201
-    return "Unauthorized", 401
+    _data.register()
+    return "Success", 201
 
 
 @pyxie.route("/unregister", methods=[C.HTTP_METHOD_DELETE])
 def unregister():
-    if _validate_api_key():
-        _data.unregister()
-        return "Success", 204
-    return "Unauthorized", 401
+    _data.unregister()
+    return "Success", 204
 
 
 @pyxie.route("/stats", methods=[C.HTTP_METHOD_GET])
 def stats():
-    if _validate_api_key():
-        stats = {}
-        for attr in dir(_data):
-            if isinstance(getattr(type(_data), attr, None), property):
-                stats[attr] = getattr(_data, attr)
-        return stats, 200
-    return "Unauthorized", 401
+    stats = {}
+    for attr in dir(_data):
+        if isinstance(getattr(type(_data), attr, None), property):
+            stats[attr] = getattr(_data, attr)
+    return stats, 200
 
 
 @pyxie.route("/metrics", methods=[C.HTTP_METHOD_GET])
 def metrics():
     # [ TODO - Issue #7] - Export prometheus formatted metrics
-    if _validate_api_key():
-        return "Metrics", 501  # Not Implemented
-    return "Unauthorized", 401
+    return "Metrics", 501  # Not Implemented
 
 
 @pyxie.route("/", methods=[C.HTTP_METHOD_GET])
@@ -58,6 +74,12 @@ def root():
     return Response(C.ONE_BY_ONE, mimetype=C.HTTP_MIME_TYPE_PNG)
 
 
+@pyxie.route("/shutdown", methods=[C.HTTP_METHOD_POST])
+def shutdown():
+    C._SHUTDOWN = True
+    return "Service is now offline. Data will be saved.", 200
+
+
 def main():
     pyxie.run(host=C.LISTEN_IP, port=C.LISTEN_PORT)