Add checksum caching (#134)

* Implement calculate_checksum on Instance

* Use new checksumming option throughout

* Rev version

Author: JBorrow

alembic/versions/38a604ac628b_add_cached_checksums.py

@@ -0,0 +1,33 @@
+# Copyright 2017 the HERA Collaboration
+# Licensed under the 2-clause BSD License.
+
+"""Add cached checksums
+
+Revision ID: 38a604ac628b
+Revises: 1def8c988372
+Create Date: 2025-02-20 11:12:54.446089
+
+"""
+import sqlalchemy as sa
+
+from alembic import op
+
+revision = "38a604ac628b"
+down_revision = "1def8c988372"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    with op.batch_alter_table("instances") as batch_op:
+        batch_op.add_column(
+            sa.Column("calculated_checksum", sa.String(), nullable=True)
+        )
+        batch_op.add_column(sa.Column("calculated_size", sa.Integer(), nullable=True))
+        batch_op.add_column(sa.Column("checksum_time", sa.DateTime(), nullable=True))
+
+
+def downgrade():
+    op.drop_column("instances", "calculated_checksum")
+    op.drop_column("instances", "calculated_size")
+    op.drop_column("instances", "checksum_time")

librarian_background/check_integrity.py

@@ -9,7 +9,7 @@
 from schedule import CancelJob
 from sqlalchemy.orm import Session
 
-from hera_librarian.utils import compare_checksums, get_hash_function_from_hash
+from hera_librarian.utils import compare_checksums
 from librarian_server.database import get_session
 from librarian_server.orm import Instance, StoreMetadata
 from librarian_server.orm.file import CorruptFile, File
@@ -81,10 +81,7 @@ def core(self, session: Session):
         for file in files:
             # Now we can check the integrity of each file.
             try:
-                hash_function = get_hash_function_from_hash(file.file.checksum)
-                path_info = store.store_manager.path_info(
-                    file.path, hash_function=hash_function
-                )
+                checksum, size = file.calculate_checksum(session=session, commit=True)
             except FileNotFoundError:
                 all_files_fine = False
                 logger.error(
@@ -98,7 +95,7 @@ def core(self, session: Session):
             # Compare checksum to database
             expected_checksum = file.file.checksum
 
-            if compare_checksums(expected_checksum, path_info.checksum):
+            if compare_checksums(expected_checksum, checksum):
                 logger.info(
                     "Instance {} on store {} has been validated (Instance: {})",
                     file.path,
@@ -120,8 +117,8 @@ def core(self, session: Session):
                 if corrupt_file is None:
                     corrupt_file = CorruptFile.new_corrupt_file(
                         instance=file,
-                        size=path_info.size,
-                        checksum=path_info.checksum,
+                        size=size,
+                        checksum=checksum,
                     )
                     session.add(corrupt_file)
                     session.commit()
@@ -134,7 +131,7 @@ def core(self, session: Session):
                     file.path,
                     store.name,
                     expected_checksum,
-                    path_info.checksum,
+                    checksum,
                     file.id,
                 )
         if all_files_fine:

librarian_background/corruption_fixer.py

@@ -17,7 +17,7 @@
     CorruptionResendResponse,
 )
 from hera_librarian.transfer import TransferStatus
-from hera_librarian.utils import compare_checksums, get_hash_function_from_hash
+from hera_librarian.utils import compare_checksums
 from librarian_server.database import get_session
 from librarian_server.orm.file import CorruptFile, File
 from librarian_server.orm.instance import Instance
@@ -91,13 +91,11 @@ def core(self, session: Session) -> bool:
 
             # Step A: Check that the file is actually corrupt
             try:
-                hash_function = get_hash_function_from_hash(potential_file.checksum)
-                store = potential_instance.store
-                path_info = store.store_manager.path_info(
-                    potential_instance.path, hash_function=hash_function
+                potential_checksum, potential_size = (
+                    potential_instance.calculate_checksum(session=session, commit=True)
                 )
 
-                if compare_checksums(potential_file.checksum, path_info.checksum):
+                if compare_checksums(potential_file.checksum, potential_checksum):
                     logger.info(
                         "CorruptFile {id} stated that file {name} was corrupt in instance {inst_id} "
                         "but we just checked the checksums: {chk_a}=={chk_b} and the file is fine "
@@ -106,7 +104,7 @@ def core(self, session: Session) -> bool:
                         name=corrupt.file_name,
                         inst_id=corrupt.instance_id,
                         chk_a=potential_file.checksum,
-                        chk_b=path_info.checksum,
+                        chk_b=potential_checksum,
                     )
                     session.delete(corrupt)
                     session.commit()

librarian_server/api/corrupt.py

@@ -13,7 +13,7 @@
     CorruptionResendRequest,
     CorruptionResendResponse,
 )
-from hera_librarian.utils import compare_checksums, get_hash_function_from_hash
+from hera_librarian.utils import compare_checksums
 from librarian_server.orm.file import File
 from librarian_server.orm.instance import Instance, RemoteInstance
 from librarian_server.orm.librarian import Librarian
@@ -99,12 +99,11 @@ def user_and_librarian_validation_flow(
             ),
         )
 
-    hash_function = get_hash_function_from_hash(file.checksum)
-    path_info = best_instance.store.store_manager.path_info(
-        best_instance.path, hash_function=hash_function
+    recomputed_checksum, _ = best_instance.calculate_checksum(
+        session=session, commit=True, force_recompute=True
     )
 
-    if not compare_checksums(file.checksum, path_info.checksum):
+    if not compare_checksums(file.checksum, recomputed_checksum):
         logger.error(
             "Our copy of the file {} is corrupt, we cannot send it to {}",
             file_name,

librarian_server/api/validate.py

@@ -27,7 +27,7 @@
     FileValidationResponse,
     FileValidationResponseItem,
 )
-from hera_librarian.utils import compare_checksums, get_hash_function_from_hash
+from hera_librarian.utils import compare_checksums
 
 from ..database import yield_session
 from ..logger import log
@@ -39,80 +39,42 @@
 
 router = APIRouter(prefix="/api/v2/validate")
 
-VALIDATION_TIMEOUT = datetime.timedelta(hours=8)
-VALIDATION_CACHE = {}
-
-
-async def cached_calculate_checksum_of_local_copy(
-    original_checksum: str,
-    original_size: int,
-    path_info_function: callable,
-    path: Path,
-    store_id: int,
-    instance_id: int,
-):
-    key = f"{original_checksum}-{instance_id}"
-
-    cached = VALIDATION_CACHE.get(key, None)
-
-    if cached is None or (
-        (datetime.datetime.now(datetime.timezone.utc) - cached[1]) > VALIDATION_TIMEOUT
-    ):
-        result = await asyncify(calculate_checksum_of_local_copy)(
-            original_checksum=original_checksum,
-            original_size=original_size,
-            path_info_function=path_info_function,
-            path=path,
-            store_id=store_id,
-            instance_id=instance_id,
-        )
-
-        VALIDATION_CACHE[key] = (result, datetime.datetime.now(datetime.timezone.utc))
-    else:
-        log.info(
-            f"Using cached result for instance {instance_id}", instance_id=instance_id
-        )
-        result = cached[0]
-
-    return result
-
 
 def calculate_checksum_of_local_copy(
     original_checksum: str,
     original_size: int,
-    path_info_function: callable,
-    path: Path,
-    store_id: int,
-    instance_id: int,
+    instance: Instance,
+    session: Session,
 ):
     start = perf_counter()
-    hash_function = get_hash_function_from_hash(original_checksum)
     try:
-        path_info = path_info_function(path, hash_function=hash_function)
+        current_checksum, current_size = instance.calculate_checksum(
+            session=session, commit=True
+        )
         response = FileValidationResponseItem(
             librarian=server_settings.name,
-            store=store_id,
-            instance_id=instance_id,
+            store=instance.store_id,
+            instance_id=instance.id,
             original_checksum=original_checksum,
             original_size=original_size,
-            current_checksum=path_info.checksum,
-            current_size=path_info.size,
+            current_checksum=current_checksum,
+            current_size=current_size,
             computed_same_checksum=compare_checksums(
-                original_checksum, path_info.checksum
+                original_checksum, current_checksum
             ),
         )
         end = perf_counter()
 
         log.debug(
-            f"Calculated path info for {instance_id} ({path_info.size} B) "
-            f"in {end - start:.2f} seconds."
+            f"Calculated path info for {response.instance_id} / {instance.path} "
+            f"({response.current_size} B) in {end - start:.2f} seconds"
         )
 
         return [response]
     except FileNotFoundError:
         # A mistakenly 'available' file that is not actually available.
         log.error(
-            f"File {path} in store {store_id} marked as available but does not exist."
+            f"File {instance.path} in store {instance.store_id} marked as available but does not exist."
         )
 
         return []
@@ -224,13 +186,11 @@ async def validate_file(
         if not instance.available:
             continue
 
-        this_checksum_info = cached_calculate_checksum_of_local_copy(
+        this_checksum_info = asyncify(calculate_checksum_of_local_copy)(
             original_checksum=file.checksum,
             original_size=file.size,
-            path_info_function=instance.store.store_manager.path_info,
-            path=instance.path,
-            store_id=instance.store.id,
-            instance_id=instance.id,
+            instance=instance,
+            session=session,
         )
 
         coroutines.append(this_checksum_info)

librarian_server/orm/instance.py

@@ -11,6 +11,7 @@
 from sqlalchemy.orm import Session
 
 from hera_librarian.deletion import DeletionPolicy
+from hera_librarian.utils import get_hash_function_from_hash
 
 from .. import database as db
 from ..settings import server_settings
@@ -50,6 +51,13 @@ class Instance(db.Base):
     available = db.Column(db.Boolean, nullable=False)
     "Whether or not this file is available on our librarian."
 
+    calculated_checksum = db.Column(db.String, nullable=True)
+    "The checksum that has been calculated for this on-disk instance"
+    calculated_size = db.Column(db.Integer, nullable=True)
+    "The size of the file that was calculated at the same time as the checksum"
+    checksum_time = db.Column(db.DateTime, nullable=True)
+    "The time at which the calculated_checksum was checked"
+
     @classmethod
     def new_instance(
         self,
@@ -124,6 +132,79 @@ def delete(
 
         return
 
+    def calculate_checksum(
+        self,
+        session: Session,
+        commit: bool = True,
+        force_recompute: bool = False,
+    ) -> tuple[str, int]:
+        """
+        Calculates the checksum of the instance on disk. It will use the stored checksum
+        in the table instead if it has not yet timed out and has been recently calculated.
+
+        Parameters
+        ----------
+        session: Session
+            The database session to use; this can be committed to if commit=True below.
+            Session must be active as we make sub-queries for file.
+        commit: bool = True
+            Whether to commit any new changes to the database.
+        force_recompute: bool = False
+            Re-compute the checksum even when a hashed version is available.
+
+        Returns
+        -------
+        checksum: str
+            The checksum, calculated or drawn from the database.
+        size: int
+            Size of the on-disk file in bytes.
+
+        Raises
+        ------
+        FileNotFoundError
+            If the file was not found on disk.
+        """
+
+        current_time = datetime.now(timezone.utc)
+
+        if (
+            (self.checksum_time is not None)
+            and (self.calculated_checksum is not None)
+            and (not force_recompute)
+        ):
+            checksum_time = self.checksum_time.astimezone(timezone.utc)
+            if (current_time - checksum_time) < server_settings.checksum_timeout:
+                logger.info(
+                    "Returning a cached checksum from {time} for instance {id} at "
+                    "{path} - {checksum}",
+                    time=checksum_time,
+                    id=self.id,
+                    path=self.path,
+                    checksum=self.calculated_checksum,
+                )
+                return self.calculated_checksum, self.calculated_size
+
+        # We must calculate the checksum fresh.
+        hash_function = get_hash_function_from_hash(self.file.checksum)
+        path_info = self.store.store_manager.path_info(self.path, hash_function)
+
+        logger.info(
+            "Calculated a fresh checksum at {time} for instance {id} at {path} - {checksum}",
+            time=current_time,
+            id=self.id,
+            path=self.path,
+            checksum=self.calculated_checksum,
+        )
+
+        self.calculated_checksum = path_info.checksum
+        self.calculated_size = path_info.size
+        self.checksum_time = current_time
+
+        if commit:
+            session.commit()
+
+        return self.calculated_checksum, self.calculated_size
+
 
 class RemoteInstance(db.Base):
     """

librarian_server/settings.py

@@ -3,6 +3,7 @@
 deserialized from the available librarian config path.
 """
 
+import datetime
 import os
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional
@@ -179,6 +180,7 @@ class ServerSettings(BaseSettings):
 
     # Checksumming options
     checksum_threads: int = 4
+    checksum_timeout: datetime.timedelta = datetime.timedelta(days=1)
 
     model_config = SettingsConfigDict(env_prefix="librarian_server_")
 

pyproject.toml

@@ -7,7 +7,7 @@ exclude=["*tests*"]
 [project]
 name="hera_librarian"
 requires-python = ">=3.10"
-version = "3.1.2"
+version = "4.0.0"
 dependencies = [
     "alembic",
     "argon2-cffi",