diff --git a/robot-server/robot_server/protocols/analysis_memcache.py b/robot-server/robot_server/protocols/analysis_memcache.py index 19280009bd5..3ba3156607f 100644 --- a/robot-server/robot_server/protocols/analysis_memcache.py +++ b/robot-server/robot_server/protocols/analysis_memcache.py @@ -63,3 +63,14 @@ def insert(self, key: K, value: V) -> None: self._pop_eldest(key) self._cache[key] = value self._cache_order.appendleft(key) + + def remove(self, key: K) -> None: + """Remove the cached element specified by the key. + + If no such element exists in cache, then simply no-op. + """ + try: + self._cache.pop(key) + self._cache_order.remove(key) # O(n) operation, use sparingly + except KeyError: + pass diff --git a/robot-server/robot_server/protocols/analysis_store.py b/robot-server/robot_server/protocols/analysis_store.py index b0ea474ec07..60ea3d8d743 100644 --- a/robot-server/robot_server/protocols/analysis_store.py +++ b/robot-server/robot_server/protocols/analysis_store.py @@ -129,8 +129,6 @@ def add_pending(self, protocol_id: str, analysis_id: str) -> AnalysisSummary: Returns: A summary of the just-added analysis. """ - # TODO (spp, 2024-03-19): cap the number of analyses being stored by - # auto-deleting old ones new_pending_analysis = self._pending_store.add( protocol_id=protocol_id, analysis_id=analysis_id ) diff --git a/robot-server/robot_server/protocols/completed_analysis_store.py b/robot-server/robot_server/protocols/completed_analysis_store.py index 58017e4398a..60780ab9cf4 100644 --- a/robot-server/robot_server/protocols/completed_analysis_store.py +++ b/robot-server/robot_server/protocols/completed_analysis_store.py @@ -21,6 +21,8 @@ _log = getLogger(__name__) +MAX_ANALYSES_TO_STORE = 5 + @dataclass class CompletedAnalysisResource: @@ -336,6 +338,7 @@ def get_ids_by_protocol(self, protocol_id: str) -> List[str]: async def add(self, completed_analysis_resource: CompletedAnalysisResource) -> None: """Add a resource to the store.""" + self._make_room_for_new_analysis(completed_analysis_resource.protocol_id) statement = analysis_table.insert().values( await completed_analysis_resource.to_sql_values() ) @@ -344,3 +347,29 @@ async def add(self, completed_analysis_resource: CompletedAnalysisResource) -> N self._memcache.insert( completed_analysis_resource.id, completed_analysis_resource ) + + def _make_room_for_new_analysis(self, protocol_id: str) -> None: + """Remove the oldest analyses in store if the number of analyses exceed the max allowed. + + Unlike protocols, protocol analysis IDs are not stored by any DB entities + other than the analysis store itself. So we do not have to worry about cleaning up + any other tables. + """ + analyses_ids = self.get_ids_by_protocol(protocol_id) + + # Delete all analyses exceeding max number allowed, + # plus an additional one to create room for the new one. + # Most existing databases will not have multiple extra analyses per protocol + # but there would be some internally that added multiple analyses before + # we started capping the number of analyses. + analyses_to_delete = analyses_ids[ + : len(analyses_ids) - MAX_ANALYSES_TO_STORE + 1 + ] + + for analysis_id in analyses_to_delete: + self._memcache.remove(analysis_id) + delete_statement = sqlalchemy.delete(analysis_table).where( + analysis_table.c.id == analysis_id + ) + with self._sql_engine.begin() as transaction: + transaction.execute(delete_statement) diff --git a/robot-server/tests/protocols/test_completed_analysis_store.py b/robot-server/tests/protocols/test_completed_analysis_store.py index f41594d0c5d..438cf8baada 100644 --- a/robot-server/tests/protocols/test_completed_analysis_store.py +++ b/robot-server/tests/protocols/test_completed_analysis_store.py @@ -2,12 +2,13 @@ import json from datetime import datetime, timezone from pathlib import Path -from typing import Optional, Dict +from typing import Optional, Dict, List import pytest from sqlalchemy.engine import Engine from decoy import Decoy +from robot_server.persistence.tables import analysis_table from robot_server.protocols.completed_analysis_store import ( CompletedAnalysisResource, CompletedAnalysisStore, @@ -261,3 +262,89 @@ async def test_get_rtp_values_and_defaults_by_analysis_from_db( decoy.when(memcache.get("analysis-id")).then_raise(KeyError()) result = await subject.get_rtp_values_and_defaults_by_analysis_id("analysis-id") assert result == resource.run_time_parameter_values_and_defaults + + +@pytest.mark.parametrize( + argnames=["existing_analysis_ids", "expected_analyses_ids_after_making_room"], + argvalues=[ + ( + [f"analysis-id-{num}" for num in range(8)], + [ + "analysis-id-4", + "analysis-id-5", + "analysis-id-6", + "analysis-id-7", + "new-analysis-id", + ], + ), + ( + [f"analysis-id-{num}" for num in range(5)], + [ + "analysis-id-1", + "analysis-id-2", + "analysis-id-3", + "analysis-id-4", + "new-analysis-id", + ], + ), + ( + [f"analysis-id-{num}" for num in range(4)], + [ + "analysis-id-0", + "analysis-id-1", + "analysis-id-2", + "analysis-id-3", + "new-analysis-id", + ], + ), + ( + [f"analysis-id-{num}" for num in range(2)], + ["analysis-id-0", "analysis-id-1", "new-analysis-id"], + ), + ([], ["new-analysis-id"]), + ], +) +async def test_add_makes_room_for_new_analysis( + subject: CompletedAnalysisStore, + memcache: MemoryCache[str, CompletedAnalysisResource], + protocol_store: ProtocolStore, + existing_analysis_ids: List[str], + expected_analyses_ids_after_making_room: List[str], + decoy: Decoy, + sql_engine: Engine, +) -> None: + """It should delete old analyses and make room for new analysis.""" + protocol_store.insert(make_dummy_protocol_resource("protocol-id")) + + # Set up the database with existing analyses + resources = [ + _completed_analysis_resource( + analysis_id=analysis_id, + protocol_id="protocol-id", + ) + for analysis_id in existing_analysis_ids + ] + for resource in resources: + statement = analysis_table.insert().values(await resource.to_sql_values()) + with sql_engine.begin() as transaction: + transaction.execute(statement) + + assert subject.get_ids_by_protocol("protocol-id") == existing_analysis_ids + await subject.add( + _completed_analysis_resource( + analysis_id="new-analysis-id", + protocol_id="protocol-id", + ) + ) + assert ( + subject.get_ids_by_protocol("protocol-id") + == expected_analyses_ids_after_making_room + ) + + removed_ids = [ + analysis_id + for analysis_id in existing_analysis_ids + if analysis_id not in expected_analyses_ids_after_making_room + ] + for analysis_id in removed_ids: + decoy.verify(memcache.remove(analysis_id)) diff --git a/robot-server/tests/protocols/test_memcache.py b/robot-server/tests/protocols/test_memcache.py index ce485d8984f..80acb184f20 100644 --- a/robot-server/tests/protocols/test_memcache.py +++ b/robot-server/tests/protocols/test_memcache.py @@ -22,3 +22,19 @@ def test_cache_retains_new_values() -> None: for val in range(1, 4): assert subject.contains(f"key-{val}") assert subject.get(f"key-{val}") == f"value-{val}" + + +def test_cache_removes_values_by_key() -> None: + """It should eject values when asked for it.""" + subject = MemoryCache(3, str, str) + for val in range(3): + subject.insert(f"key-{val}", f"value-{val}") + subject.remove("key-1") + assert not subject.contains("key-1") + + # Make sure cache order is updated + assert subject.contains("key-0") and subject.contains("key-2") + subject.insert("key-4", "value-4") + assert subject.contains("key-0") + subject.insert("key-5", "value-5") + assert not subject.contains("key-0")