# # This file is distributed under the MIT License. See LICENSE.md for details. # from __future__ import annotations from collections import defaultdict from collections.abc import Buffer from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime from pathlib import Path from typing import AsyncGenerator, Iterable, Mapping from revng.pypeline import __version__ as version from revng.pypeline.container import ContainerID from revng.pypeline.model import Model, ModelPathSet from revng.pypeline.object import ObjectID from revng.pypeline.task.pipe import ObjectDependencies, PipeCustomInvalidation from revng.pypeline.utils import Locked from revng.pypeline.utils.registry import get_singleton from .file_provider import FileRequest from .storage_provider import ConfigurationId, ContainerLocation, FileStorageEntry from .storage_provider import InvalidatedObjects, ObjectsToInvalidate, ProjectID, ProjectMetadata from .storage_provider import SavepointID, SavePointsRange, StorageProvider from .storage_provider import StorageProviderFactory from .util import check_kind_structure, compute_hash @dataclass(frozen=True) class DependencyEntry: savepoint_start: SavepointID savepoint_end: SavepointID container_id: ContainerID configuration_id: ConfigurationId object_id: ObjectID class InMemoryStorageProviderFactory(StorageProviderFactory): def __init__(self, url: str): assert url == "memory://" self.providers: Locked[dict[ProjectID | None, Locked[InMemoryStorageProvider]]] = Locked({}) @classmethod def scheme(cls) -> str: return "memory" @asynccontextmanager async def get( self, base_directory: Path, project_id: ProjectID | None, token: str | None, cache_dir: str | None, ) -> AsyncGenerator[StorageProvider]: async with self.providers() as providers: project_provider: Locked[InMemoryStorageProvider] | None = providers.get(project_id) if project_provider is None: project_provider = Locked(InMemoryStorageProvider()) providers[project_id] = project_provider # Release the global lock and acquire the project-specific one so other # projects can proceed in parallel async with project_provider() as project_data: yield project_data class InMemoryStorageProvider(StorageProvider): """A simple in-memory storage provider for testing purposes. This is not thread-safe and should not be used in production. """ def __init__(self): check_kind_structure() self.model: Model = get_singleton(Model)() self.storage: dict[ContainerLocation, dict[ObjectID, bytes]] = {} self.dependencies: dict[str, list[DependencyEntry]] = defaultdict(list) self.last_change = datetime.now() self.files: dict[str, bytes] = {} self.epoch = 0 self.custom_dependencies: dict[tuple[int, str], PipeCustomInvalidation] = {} def has( self, location: ContainerLocation, keys: Iterable[ObjectID], ) -> Iterable[ObjectID]: if location not in self.storage: return [] storage = self.storage[location] return [key for key in keys if key in storage] def get( self, location: ContainerLocation, keys: Iterable[ObjectID], ) -> Mapping[ObjectID, bytes]: if location not in self.storage: raise KeyError(f"Savepoint {location} not found in storage.") storage = self.storage[location] return {k: storage[k] for k in keys if k in storage} def add_dependencies( self, savepoint_range: SavePointsRange, configuration_id: ConfigurationId, deps: ObjectDependencies, ) -> None: for container_name, obj, path in deps: self.dependencies[path].append( DependencyEntry( savepoint_range.start, savepoint_range.end, container_name, configuration_id, obj, ) ) self.last_change = datetime.now() def put( self, location: ContainerLocation, values: Mapping[ObjectID, Buffer], ) -> None: self.storage.setdefault(location, {}) for key, value in values.items(): self.storage[location][key] = bytes(value) self.last_change = datetime.now() def invalidate( self, invalidation_list: ModelPathSet, additional_objects: list[ObjectsToInvalidate] ) -> InvalidatedObjects: invalidated: InvalidatedObjects = defaultdict(set) # Set of entries that will be collected from self.dependencies object_to_delete: set[DependencyEntry] = set() # Keys that exist on self.dependencies that will be deleted paths_to_delete: list[str] = [] # Append the additional objects to delete for object_set in additional_objects: for object_ in object_set.objects: object_to_delete.add( DependencyEntry( object_set.savepoint_range.start, object_set.savepoint_range.end, object_set.container_id, object_set.configuration_id, object_, ) ) # Retrieve entries from self.dependencies that match the provided paths for path in invalidation_list: if path in self.dependencies: paths_to_delete.append(path) object_to_delete.update(self.dependencies[path]) # For each DependencyEntry, find the matching entries in self.storage # TODO: this double loop is very inefficient for entry in object_to_delete: for container_loc, objects in self.storage.items(): # For an entry to match, the savepoint must be in the savepoint # range and the configuration_id must match if ( container_loc.savepoint_id < entry.savepoint_start or container_loc.savepoint_id > entry.savepoint_end or container_loc.configuration_id != entry.configuration_id ): continue # Check that the object is related (the same object, a parent # or a child) to the invalidated object for object_ in objects: if entry.object_id.is_related(object_): invalidated[container_loc].add(object_) # Actually delete the data for container_loc, objects_to_delete in invalidated.items(): self.storage[container_loc] = { k: v for k, v in self.storage[container_loc].items() if k not in objects_to_delete } # Delete the paths that were previously found for path in paths_to_delete: del self.dependencies[path] self.last_change = datetime.now() return dict(invalidated) def get_epoch(self) -> int: return self.epoch def get_model(self) -> tuple[Model, int]: return (self.model.clone(), self.epoch) def set_model(self, new_model: Model) -> int: if self.model != new_model: self.epoch += 1 self.model = new_model.clone() self.last_change = datetime.now() return self.epoch def metadata(self) -> ProjectMetadata: """ Fetch metadata about the current project """ return ProjectMetadata( last_change=self.last_change, version=version, ) def prune_objects(self): """ Prunes all the objects (except metadata) from storage """ self.storage.clear() self.dependencies.clear() def put_files_in_storage(self, files: list[FileStorageEntry]) -> list[str]: result = [] for file in files: contents = None if file.contents is not None: contents = file.contents elif file.path is not None: contents = file.path.read_bytes() assert contents is not None hash_ = compute_hash(contents) self.files[hash_] = contents result.append(hash_) return result def get_files_from_storage(self, requests: list[FileRequest]) -> dict[str, bytes]: return {r.hash: self.files[r.hash] for r in requests} def add_custom_invalidation_data( self, pipe_id: int, configuration_hash: str, data: PipeCustomInvalidation ): storage = [[(oid, bytes(d)) for oid, d in entry] for entry in data] self.custom_dependencies[(pipe_id, configuration_hash)] = storage def get_custom_invalidation_data( self, pipe_id: int, configuration_hash: str ) -> PipeCustomInvalidation: return self.custom_dependencies[(pipe_id, configuration_hash)]