feat: Added missing/stale worker threshold (#14)

.env_template

@@ -19,7 +19,11 @@ POSTGRES_URI="postgresql://pictrsafety:ThatSecretPGSQLPass@127.0.0.1/pictrsafety
 # Note than 127.0.0.1 is always allowed.
 #KNOWN_PICTRS_IPS='["192.168.0.1"]'
 # If this amount of images are waiting for scanning, start retuning "OK" for each new scan after that
-# This will prevent pict-rs from rejecting each and every upload because your scanning workers crashed or something
+# This will prevent pict-rs from rejecting each and every upload because your scanning workers are too slow or crashed or something
 # and the admin is AFK
 # Set this value to 0 to never bypass the scan.
-SCAN_BYPASS_THRESHOLD=10
+SCAN_BYPASS_THRESHOLD=10
+# If this amount of seconds have passed since a scanning worker last checked in, start returning "OK" for each new scan after that
+# This will prevent pict-rs from rejecting each and every upload because your scanning workers have gone missing
+# Set this value to 0 to never bypass the scan.
+MISSING_WORKER_THRESHOLD=5

pictrs_safety_api/apis/v1/base.py

@@ -13,6 +13,7 @@
 from pictrs_safety_api.database import functions as database
 from pictrs_safety_api import exceptions as e
 from pictrs_safety_api import enums
+from pictrs_safety_api.classes.worker import worker
 
 api = Namespace('v1', 'API Version 1' )
 
@@ -51,9 +52,13 @@ def post(self,pictrs_id):
                 raise e.Unauthorized("You are not authorized to use this service", f"Unauthorized IP: {request.remote_addr}")
         elif pictrs_id not in json.loads(os.getenv("KNOWN_PICTRS_IDS", "[]")):
             raise e.Unauthorized("You are not authorized to use this service", f"Unauthorized ID: {pictrs_id}")
+        if worker.is_stale():
+            logger.warning(f"Returning OK due to stale worker. Last seen: {worker.last_check_in}")
+            return {"message": "Worker Stale"}, 200 
         scan_threshold = int(os.getenv("SCAN_BYPASS_THRESHOLD", 10))
         if scan_threshold > 0 and database.count_waiting_scan_requests() > scan_threshold:
-            return {"message": "Image OK"}, 200 
+            logger.warning(f"Returning OK due to full queue. Queue: {database.count_waiting_scan_requests()}")
+            return {"message": "Queue Full"}, 200 
         self.args = self.post_parser.parse_args()
         file = self.args["file"]
         if not file:
@@ -144,10 +149,11 @@ def get(self):
         '''Pick up an image to safety validate
         '''
         # I don't get why this is not using the import from earlier...
-        logger.debug(request.remote_addr)
+        # logger.debug(request.remote_addr)
         self.args = self.get_parser.parse_args()
         if os.getenv("FEDIVERSE_SAFETY_WORKER_AUTH") != self.args.apikey:
             raise e.Forbidden("Access Denied")
+        worker.check_in()
         pop: ScanRequest = database.find_waiting_scan_request()
         if not pop:
             return {"message": "Nothing to do"},204

pictrs_safety_api/classes/request.py

@@ -1,6 +1,5 @@
-import dateutil.relativedelta
 from datetime import datetime
-from sqlalchemy import Enum, UniqueConstraint
+from sqlalchemy import Enum
 
 from loguru import logger
 from pictrs_safety_api.flask import db

pictrs_safety_api/classes/worker.py

@@ -0,0 +1,20 @@
+import os
+from datetime import datetime, timedelta
+
+class Worker():
+    last_check_in = None
+
+    def __init__(self):
+        self.check_in()
+
+    def check_in(self):
+        self.last_check_in = datetime.utcnow()
+
+    def is_stale(self):
+        threshold = int(os.getenv("MISSING_WORKER_THRESHOLD", 5))
+        # When threshold is 0, the worker is never stale
+        if threshold == 0:
+            return False
+        return self.last_check_in < datetime.utcnow() - timedelta(seconds=threshold)
+
+worker = Worker()

pictrs_safety_api/database/functions.py

@@ -1,8 +1,4 @@
-import time
-import uuid
-import json
 from loguru import logger
-from datetime import datetime, timedelta
 from sqlalchemy import func, or_, and_, not_, Boolean
 from sqlalchemy.orm import noload
 from pictrs_safety_api.flask import db