Cleanup

Author: the-glu

monitoring/mock_uss/ridsp/database.py

@@ -1,13 +1,20 @@
 import json
+from datetime import timedelta
 
+import arrow
 from implicitdict import ImplicitDict, Optional
 
+from monitoring.mock_uss.app import webapp
 from monitoring.monitorlib.multiprocessing import SynchronizedValue
 from monitoring.monitorlib.rid_automated_testing import injection_api
 
 from .behavior import ServiceProviderBehavior
 from .user_notifications import ServiceProviderUserNotifications
 
+DB_CLEANUP_INTERVAL = timedelta(hours=1)
+FLIGHTS_LIMIT = timedelta(hours=1)
+NOTIFICATIONS_LIMIT = timedelta(hours=1)
+
 
 class TestRecord(ImplicitDict):
     """Representation of RID SP's record of a set of injected test flights"""
@@ -25,6 +32,14 @@ def __init__(self, **kwargs):
 
         super().__init__(**kwargs)
 
+    def cleanup_flights(self):
+        self.flights = [
+            flight
+            for flight in self.flights
+            if flight.get_span()[1]
+            and flight.get_span()[1] + FLIGHTS_LIMIT > arrow.utcnow().datetime  # pyright: ignore[reportOptionalOperand]
+        ]
+
 
 class Database(ImplicitDict):
     """Simple pseudo-database structure tracking the state of the mock system"""
@@ -33,8 +48,36 @@ class Database(ImplicitDict):
     behavior: ServiceProviderBehavior = ServiceProviderBehavior()
     notifications: ServiceProviderUserNotifications = ServiceProviderUserNotifications()
 
+    def cleanup_notifications(self):
+        self.notifications.cleanup(NOTIFICATIONS_LIMIT)
+
+    def cleanup_flights(self):
+        to_cleanup = []
+
+        for test_id, test in self.tests.items():
+            if test.flights:
+                test.cleanup_flights()
+
+                if not test.flights:
+                    to_cleanup.append(test_id)
+
+        for test_id in to_cleanup:
+            del self.tests[test_id]
+
 
 db = SynchronizedValue[Database](
     Database(),
     decoder=lambda b: ImplicitDict.parse(json.loads(b.decode("utf-8")), Database),
 )
+
+TASK_DATABASE_CLEANUP = "ridsp database cleanup"
+
+
+@webapp.periodic_task(TASK_DATABASE_CLEANUP)
+def database_cleanup() -> None:
+    with db.transact() as tx:
+        tx.value.cleanup_notifications()
+        tx.value.cleanup_flights()
+
+
+webapp.set_task_period(TASK_DATABASE_CLEANUP, DB_CLEANUP_INTERVAL)

monitoring/mock_uss/ridsp/user_notifications.py

@@ -17,8 +17,6 @@
 
 from . import database
 
-NOTIFICATIONS_LIMIT = datetime.timedelta(hours=1)
-
 
 class ServiceProviderUserNotifications(ImplicitDict):
     user_notifications: list[UserNotification] = []
@@ -48,13 +46,11 @@ def create_notifications_if_needed(self, record: "database.TestRecord"):
         ):
             self.record_notification(message=notif_str, observed_at=notif_date)
 
-        self.cleanup()
-
-    def cleanup(self):
+    def cleanup(self, limit):
         self.user_notifications = [
             notif
             for notif in self.user_notifications
-            if notif.observed_at + NOTIFICATIONS_LIMIT > arrow.utcnow().datetime
+            if notif.observed_at.value.datetime + limit > arrow.utcnow().datetime
         ]