Fix load shedding error handling and cleanup

Author: corlettb

app/errors.py

@@ -56,7 +56,7 @@ def authentication_error(error: AuthError) -> ResponseReturnValue:
 
     @blueprint.errorhandler(ServiceUnavailableError)
     def service_unavailable_error(error: ServiceUnavailableError) -> ResponseReturnValue:
-        response = jsonify(error.to_dict_v2())
+        response = jsonify(result="error", message=error.message)
         response.status_code = 429
         response.headers["Retry-After"] = str(error.retry_after)
         current_app.logger.info(error)

app/load_shedding.py

@@ -52,6 +52,8 @@ def __init__(self, window_seconds: float = 60.0) -> None:
         """
         self._requests: dict[str, deque[float]] = {}
         self._window = window_seconds
+        self._last_cleanup = 0.0
+        self._cleanup_interval = window_seconds
 
     def track_request(self, service_id: str) -> None:
         """Record a request from this service.
@@ -68,6 +70,11 @@ def track_request(self, service_id: str) -> None:
         # Lazy cleanup of old entries to keep memory bounded
         self._cleanup_old_entries(service_id, now)
 
+        # Periodic cleanup across all services to avoid unbounded growth
+        if now - self._last_cleanup >= self._cleanup_interval:
+            self._cleanup_all(now)
+            self._last_cleanup = now
+
     def _cleanup_old_entries(self, service_id: str, now: float) -> None:
         """Remove entries older than the window.
 
@@ -84,6 +91,22 @@ def _cleanup_old_entries(self, service_id: str, now: float) -> None:
         if not queue:
             del self._requests[service_id]
 
+    def _cleanup_all(self, now: float) -> None:
+        """Remove expired entries for all services.
+
+        Args:
+            now: Current timestamp
+        """
+        cutoff = now - self._window
+
+        for service_id in list(self._requests.keys()):
+            queue = self._requests[service_id]
+            while queue and queue[0] < cutoff:
+                queue.popleft()
+
+            if not queue:
+                del self._requests[service_id]
+
     def get_volumes(self) -> dict[str, int]:
         """Get current volumes for all services that have recent requests.
 

app/v2/errors.py

@@ -10,6 +10,7 @@
 from app.authentication.auth import AuthError
 from app.clients.document_download import DocumentDownloadError
 from app.errors import InvalidRequest
+from app.load_shedding import ServiceUnavailableError
 
 
 class TooManyRequestsError(InvalidRequest):
@@ -84,7 +85,23 @@ def to_dict_v2(self):
         }
 
 
-def register_errors(blueprint):
+def _register_rate_limit_errors(blueprint):
+    @blueprint.errorhandler(ServiceUnavailableError)
+    def service_unavailable_error(error: ServiceUnavailableError):
+        current_app.logger.info(error)
+        response = jsonify(error.to_dict_v2())
+        response.status_code = 429
+        response.headers["Retry-After"] = str(error.retry_after)
+        return response
+
+    @blueprint.errorhandler(InvalidRequest)
+    def invalid_data(error):
+        current_app.logger.info(error)
+        response = jsonify(error.to_dict_v2()), error.status_code
+        return response
+
+
+def _register_validation_errors(blueprint):
     @blueprint.errorhandler(InvalidRecipientError)
     def invalid_format(error):
         current_app.logger.info(error)
@@ -109,36 +126,36 @@ def invalid_phone(error):
             400,
         )
 
-    @blueprint.errorhandler(InvalidRequest)
-    def invalid_data(error):
-        current_app.logger.info(error)
-        response = jsonify(error.to_dict_v2()), error.status_code
-        return response
-
-    @blueprint.errorhandler(DocumentDownloadError)
-    def document_download_error(error: DocumentDownloadError):
-        current_app.logger.info(error)
-        # cast it to a BadRequestError to ensure we format the error well
-        bad_request_exc = BadRequestError(message=error.message)
-        response = jsonify(bad_request_exc.to_dict_v2()), error.status_code
-        return response
-
     @blueprint.errorhandler(JsonSchemaValidationError)
     def validation_error(error):
         current_app.logger.info(error)
         return jsonify(json.loads(error.message)), 400
 
+
+def _register_data_errors(blueprint):
     @blueprint.errorhandler(NoResultFound)
     @blueprint.errorhandler(DataError)
     def no_result_found(e):
         current_app.logger.info(e, exc_info=True)
         return jsonify(status_code=404, errors=[{"error": e.__class__.__name__, "message": "No result found"}]), 404
 
+
+def _register_auth_errors(blueprint):
     @blueprint.errorhandler(AuthError)
     def auth_error(error):
         current_app.logger.info("API AuthError, client: %s error: %s", request.headers.get("User-Agent"), error)
         return jsonify(error.to_dict_v2()), error.code
 
+
+def _register_misc_errors(blueprint):
+    @blueprint.errorhandler(DocumentDownloadError)
+    def document_download_error(error: DocumentDownloadError):
+        current_app.logger.info(error)
+        # cast it to a BadRequestError to ensure we format the error well
+        bad_request_exc = BadRequestError(message=error.message)
+        response = jsonify(bad_request_exc.to_dict_v2()), error.status_code
+        return response
+
     @blueprint.errorhandler(EventletTimeout)
     def eventlet_timeout(error):
         current_app.logger.exception(error)
@@ -156,3 +173,11 @@ def internal_server_error(error):
             jsonify(status_code=500, errors=[{"error": error.__class__.__name__, "message": "Internal server error"}]),
             500,
         )
+
+
+def register_errors(blueprint):
+    _register_rate_limit_errors(blueprint)
+    _register_validation_errors(blueprint)
+    _register_data_errors(blueprint)
+    _register_auth_errors(blueprint)
+    _register_misc_errors(blueprint)