Merge pull request #29 from dev-dull/28-clean-shutdown

#28 - implement a safe shutdown endpoint

Author: dev-dull

README.md

@@ -167,3 +167,11 @@ Note that unregistering an ID is destructive and all data for that `id` will be
 user@shell> curl -Ss -X DELETE -H 'X-Api-Key: your-api-key-here' 'http://127.0.0.1:5000/unregister?id=testing'
 Success
 ```
+
+### Cleanly shutting down PyXIE
+If PyXIE is killed while in the middle of persisting data to disk, it will likely result in a corrupted file. To ensure that your data is complete and well formed, use one of your defined API keys ot send a `POST` request to the `/shutdown` endpoint like in the following example. This will tell PyXIE to write all data to disk and exit.
+
+```bash
+user@shell>  curl -Ss -X POST -H 'X-Api-Key: your-api-key-here' 'http://127.0.0.1:5000/shutdown'
+curl: (52) Empty reply from server
+```

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,7 @@
+import os
 from ddb import DDB
-from time import time
 from constfig import C
-from flask import Flask, Response, request
+from flask import Flask, Response, request, abort
 
 
 pyxie = Flask(C.APP_NAME)
@@ -13,39 +13,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 +73,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)