mirror of
https://github.com/revng/revng
synced 2026-06-21 14:07:57 +00:00
Giacomo Vercesi
2f0f320711
Fix the `revng2 project init` command to work with all `StorageProvider`s, not just the local one. It still specializes for the local one to allow creating the model file.
368 lines
12 KiB
Python
368 lines
12 KiB
Python
#
|
|
# This file is distributed under the MIT License. See LICENSE.md for details.
|
|
#
|
|
|
|
from __future__ import annotations
|
|
|
|
import enum
|
|
import json
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Buffer
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Annotated, AsyncContextManager, Collection, Iterable, Mapping
|
|
from urllib.parse import urlparse
|
|
|
|
from revng.pypeline.container import ConfigurationId, ContainerID
|
|
from revng.pypeline.model import Model, ModelPathSet
|
|
from revng.pypeline.object import ObjectID, ObjectSet
|
|
from revng.pypeline.storage.notification_queue import LOCAL_QUEUE
|
|
from revng.pypeline.task.pipe import ObjectDependencies, PipeCustomInvalidation
|
|
from revng.pypeline.utils.registry import get_registry
|
|
|
|
from .file_provider import FileProvider, FileRequest
|
|
|
|
if TYPE_CHECKING:
|
|
from revng.pypeline.pipeline import Pipeline
|
|
|
|
SavepointID = Annotated[
|
|
int,
|
|
"""
|
|
An integer that represents a savepoint, we want this to be an integer so we
|
|
can assign them doing a DFS traversal of the pipeline, which in turns allows
|
|
to efficiently represent a subtree of savepoints as a continuous range of
|
|
integers to avoid storing the dependencies multiple times.
|
|
""",
|
|
]
|
|
|
|
ProjectID = Annotated[
|
|
str,
|
|
"""
|
|
An unique identifier for a project.
|
|
""",
|
|
]
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ContainerLocation:
|
|
savepoint_id: SavepointID
|
|
container_id: ContainerID
|
|
configuration_id: ConfigurationId
|
|
|
|
|
|
InvalidatedObjects = dict[ContainerLocation, set[ObjectID]]
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ProjectMetadata:
|
|
# The version string of the last client that used the StorageProvider
|
|
version: str
|
|
# The hash of the pipeline-description used last
|
|
pipeline_description_hash: str
|
|
|
|
# Last time data was fetched from the StorageProvider
|
|
last_fetch: datetime
|
|
# Last time objects were saved in the StorageProvider
|
|
last_object_save: datetime
|
|
# Last time the model was changed in the StorageProvider
|
|
last_model_save: datetime
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class SavePointsRange:
|
|
"""A range of savepoints, with inclusive extremes.
|
|
When the savepoints IDs are assigned in a DFS traversal of the pipeline,
|
|
any subtree of savepoints can be represented as a continuous range of integers,
|
|
avoiding the need to store the dependencies multiple times.
|
|
The starts are assigned with a pre-order traversal, and the ends are assigned
|
|
with a post-order traversal. The order is increase only when a savepoint is
|
|
visited for the first time.
|
|
"""
|
|
|
|
start: SavepointID
|
|
"""
|
|
This is the smallest savepoint ID of the subtree of a node (root included).
|
|
Thus, if the node is a SavePoint, this is the ID of the SavePoint itself.
|
|
"""
|
|
|
|
end: SavepointID
|
|
"""This is the largest savepoint ID of the subtree of a node (root included).
|
|
Thus, if the node is a SavePoint, this is start + the number of savepoints
|
|
present in the subtree - 1 (minus the root itself)."""
|
|
|
|
def __contains__(self, item: object) -> bool:
|
|
if not isinstance(item, int):
|
|
return False
|
|
return self.start <= item <= self.end
|
|
|
|
def __len__(self) -> int:
|
|
"""The number of savepoints in the range."""
|
|
return self.end - self.start + 1
|
|
|
|
|
|
@dataclass
|
|
class FileStorageEntry:
|
|
# Name to use for the file in case it needs to be stored on disk
|
|
name: str
|
|
# Path of the file to store, this will avoid loading the file contents in
|
|
# memory all at once
|
|
path: Path | None = field(default=None)
|
|
# Contents of the file to store
|
|
contents: bytes | None = field(default=None)
|
|
|
|
def __post_init__(self):
|
|
assert int(self.path is not None) + int(self.contents is not None) == 1
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ObjectsToInvalidate:
|
|
savepoint_range: SavePointsRange
|
|
container_id: ContainerID
|
|
configuration_id: ConfigurationId
|
|
objects: ObjectSet
|
|
|
|
|
|
@dataclass
|
|
class SetModelResult:
|
|
# The new epoch after changing the model
|
|
epoch: int
|
|
# The objects that have been pruned from storage, stored as a dictionary
|
|
# that maps the container location to a list of objects
|
|
invalidated_objects: InvalidatedObjects
|
|
|
|
|
|
@dataclass
|
|
class PipeDependencies:
|
|
pipe_id: int
|
|
savepoints_range: SavePointsRange
|
|
configuration: ConfigurationId
|
|
dependencies: ObjectDependencies
|
|
custom_invalidation: PipeCustomInvalidation
|
|
|
|
def empty(self) -> bool:
|
|
return len(self.dependencies) == 0 and not self.has_custom_invalidation()
|
|
|
|
def has_custom_invalidation(self):
|
|
return not all(len(x) == 0 for x in self.custom_invalidation)
|
|
|
|
|
|
class LockType(enum.Enum):
|
|
ARTIFACT = enum.auto()
|
|
ANALYSIS = enum.auto()
|
|
|
|
|
|
class StorageProviderFactory(ABC):
|
|
"""
|
|
A possibly stateful factory for a specific storage provider.
|
|
This abstraction is needed to optimize daemon, as the cloud version needs to
|
|
be stateless and instantiate a new storage provider client each time, while
|
|
the local version can reuse the same client.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def __init__(self, url: str):
|
|
"""
|
|
The cli will decide how to instantiate the factory using a url argument
|
|
so this url has to contain all the information needed to instantiate the
|
|
factory.
|
|
"""
|
|
|
|
@classmethod
|
|
@abstractmethod
|
|
def scheme(cls) -> str:
|
|
"""
|
|
The scheme that identifies this storage provider, i.e. the initial part
|
|
of the URL before the "://".
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get(
|
|
self,
|
|
base_directory: Path,
|
|
pipeline: Pipeline,
|
|
lock_type: LockType,
|
|
project_id: ProjectID | None,
|
|
token: str | None,
|
|
cache_dir: str | None,
|
|
) -> AsyncContextManager[StorageProvider]:
|
|
"""
|
|
Get a storage provider instance for the given project.
|
|
This method has to receive all the arguments needed by all possible
|
|
implementations, so it's common that the implementation will ignore some
|
|
of them. The method returns an async context manager that yields the
|
|
storage provider. This is done because we need to lock the storage provider
|
|
based on the project ID because only one user at time can modify a project.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_notification_websocket(self) -> str | None:
|
|
"""
|
|
Get the URL of the notification websocket. If the function returns None
|
|
then the internal websocket will be used.
|
|
"""
|
|
|
|
def init(self, directory: Path) -> bool:
|
|
"""
|
|
Initialize the specified directory to become a storage directory, can
|
|
be overridden by local storage providers. Returns true if the directory
|
|
was initted or false if it was previously initted.
|
|
"""
|
|
return True
|
|
|
|
|
|
# TODO: find a more suitable name
|
|
def storage_provider_factory_factory(url: str) -> StorageProviderFactory:
|
|
"""Create a storage provider factory from an url."""
|
|
scheme = urlparse(url).scheme
|
|
factories = get_registry(StorageProviderFactory) # type: ignore [type-abstract]
|
|
for factory_type in factories.values():
|
|
if factory_type.scheme() == scheme:
|
|
return factory_type(url)
|
|
available_schemes = ", ".join(factories.keys())
|
|
raise ValueError(
|
|
f"Unknown storage provider scheme {scheme}."
|
|
f" The available schemes are: {available_schemes}."
|
|
)
|
|
|
|
|
|
class StorageProvider(ABC):
|
|
"""This is the general interface for something that caches containers.
|
|
This can be in memory, on disk, in a database, etc.
|
|
This is a singleton and there should never be more than one instance of it.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def has(
|
|
self,
|
|
location: ContainerLocation,
|
|
keys: Collection[ObjectID],
|
|
) -> Iterable[ObjectID]:
|
|
"""
|
|
Get the available objects from the storage.
|
|
If the object is not found, it will not be included in the result.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get(
|
|
self,
|
|
location: ContainerLocation,
|
|
keys: Collection[ObjectID],
|
|
) -> Mapping[ObjectID, bytes]:
|
|
"""
|
|
For each objects, return bytes that the container can ingest to
|
|
deserialize the object.
|
|
All the object **have to** be present, as one should call `get_available`
|
|
first to check it.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def add_objects(
|
|
self,
|
|
dependencies: list[PipeDependencies],
|
|
objects: Mapping[ContainerLocation, Mapping[ObjectID, Buffer]],
|
|
) -> None:
|
|
"""
|
|
Store dependencies and custom invalidation for an arbitrary number of
|
|
pipes while storing objects to storage. This guarantees that the
|
|
objects stored and their dependencies are stored atomically and avoids
|
|
dangling dependencies.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_epoch(self) -> int:
|
|
"""Get the epoch, i.e. the model version number."""
|
|
|
|
@abstractmethod
|
|
def get_model(self) -> tuple[Model, int]:
|
|
"""Get the model and the epoch."""
|
|
|
|
@abstractmethod
|
|
def set_model(
|
|
self,
|
|
new_model: Model,
|
|
changed_paths: ModelPathSet,
|
|
custom_invalidations: list[ObjectsToInvalidate],
|
|
) -> SetModelResult:
|
|
"""
|
|
Inform the storage that the model has changed, with the list of changed
|
|
model paths and object to explicitly invalidate (from custom
|
|
invalidation). This function will compute the overall list of objects
|
|
to invalidate and set the model. It will return the new epoch (which
|
|
will be current epoch + 1 if the model changed, or current epoch if it
|
|
didn't) and the exhaustive list of objects that have been deleted due
|
|
to the changes in the model.
|
|
This requires the new model and the list of changed paths separately so
|
|
that the storage provider does not need to run operations (diff/apply)
|
|
on the model.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def metadata(self) -> ProjectMetadata:
|
|
"""
|
|
Fetch metadata about the current project
|
|
"""
|
|
|
|
@abstractmethod
|
|
def prune_objects(self):
|
|
"""
|
|
Prunes all the objects (except metadata) from storage. This is a debug
|
|
method and should not be relied upon in other parts of the codebase.
|
|
This method can throw an exception if the provider does not support
|
|
pruning.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def put_files_in_storage(self, files: list[FileStorageEntry]) -> list[str]:
|
|
"""
|
|
Put a file in storage, which can be later retrieved with the
|
|
`get_files_from_storage` method. This method returns a list of
|
|
hashes of the submitted files.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_files_from_storage(self, requests: list[FileRequest]) -> dict[str, bytes]:
|
|
"""
|
|
Get a file from storage, given a list of requests. Returns a dictionary
|
|
mapping the hash to the contents.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_custom_invalidation_data(
|
|
self, pipe_id: int, configuration_hash: str
|
|
) -> PipeCustomInvalidation:
|
|
"""
|
|
Retrieve custom invalidation data previously stored by
|
|
`add_custom_invalidation_data`.
|
|
"""
|
|
|
|
@staticmethod
|
|
def _send_local_invalidation(invalidated: InvalidatedObjects, epoch: int):
|
|
"""
|
|
Send invalidation data to the local queue. Only to be used on local
|
|
storage providers where notifications are not handled by the provider.
|
|
"""
|
|
|
|
payload = {
|
|
"type": "invalidation",
|
|
"epoch": epoch,
|
|
"invalidated": [
|
|
{
|
|
"savepoint_id": location.savepoint_id,
|
|
"container_id": location.container_id,
|
|
"configuration": location.configuration_id,
|
|
"object_ids": [object_id.serialize() for object_id in object_ids],
|
|
}
|
|
for location, object_ids in invalidated.items()
|
|
],
|
|
}
|
|
LOCAL_QUEUE.send(json.dumps(payload).encode())
|
|
|
|
|
|
class StorageProviderFileProvider(FileProvider):
|
|
def __init__(self, provider: StorageProvider):
|
|
self._provider = provider
|
|
|
|
def get_files(self, requests: list[FileRequest]) -> dict[str, bytes]:
|
|
return self._provider.get_files_from_storage(requests)
|