mirror of
https://github.com/revng/revng
synced 2026-06-21 14:07:57 +00:00
Giacomo Vercesi
6a88dc3d6a
Implement a downstream load/save handler in `revng.internal.api.Manager` class which handles loading and saving files from a local directory to a s3 server.
336 lines
12 KiB
Python
336 lines
12 KiB
Python
#
|
|
# This file is distributed under the MIT License. See LICENSE.md for details.
|
|
#
|
|
|
|
import os
|
|
import re
|
|
import sys
|
|
import traceback
|
|
from concurrent.futures import Future, ThreadPoolExecutor
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from tempfile import NamedTemporaryFile
|
|
from threading import Event, Lock
|
|
from typing import ParamSpec, Protocol, Sequence, TypeVar, cast
|
|
from uuid import uuid4
|
|
|
|
import boto3
|
|
import botocore
|
|
import yaml
|
|
|
|
|
|
class Synchronizer(Protocol):
|
|
"""A Synchronizer is a class that takes care of saving the manager's
|
|
workdir somewhere other than the local disk. This class needs to be
|
|
thread-safe, especially for the `save` method, since it is expected that
|
|
multiple threads will call it."""
|
|
def load(self, path: Path) -> bool:
|
|
...
|
|
|
|
def save(self, path: Path) -> bool:
|
|
...
|
|
|
|
def set_credentials(self, credentials: str):
|
|
...
|
|
|
|
|
|
class NoopSynchronizer(Synchronizer):
|
|
def load(self, path: Path) -> bool:
|
|
return True
|
|
|
|
def save(self, path: Path) -> bool:
|
|
return True
|
|
|
|
def set_credentials(self, credentials: str):
|
|
pass
|
|
|
|
|
|
executor = ThreadPoolExecutor(8)
|
|
|
|
T = TypeVar("T")
|
|
P = ParamSpec("P")
|
|
|
|
|
|
class BotocoreSession(botocore.session.Session):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.credentials = botocore.credentials.Credentials("", "")
|
|
|
|
def get_credentials(self):
|
|
return self.credentials
|
|
|
|
def set_credentials(self, access_key: str, secret_key: str): # type: ignore
|
|
self.credentials.access_key = access_key
|
|
self.credentials.secret_key = secret_key
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FileData:
|
|
s3_filename: str
|
|
index: int = field(default=0)
|
|
encoding: str = field(default="None")
|
|
|
|
|
|
@contextmanager
|
|
def with_event(ev: Event):
|
|
"""Exception/return-safe way of set()-ing an Event. When the __exit__
|
|
method of the context manager is called the event will be set"""
|
|
yield None
|
|
ev.set()
|
|
|
|
|
|
def gather_bools(future_list: Sequence[Future[bool]], exception_message: str) -> bool:
|
|
"""Similar to asyncio.gather for a list of Future[bool]. Will gracefully
|
|
manage multiple exceptions by returning an ExceptionGroup. The returned
|
|
result is True if all the awaited booleans are true and false otherwise."""
|
|
exceptions: list[Exception] = []
|
|
result = True
|
|
for future in future_list:
|
|
try:
|
|
if not future.result():
|
|
result = False
|
|
except Exception as ex:
|
|
exceptions.append(ex)
|
|
|
|
if len(exceptions) > 0:
|
|
raise ExceptionGroup(exception_message, exceptions)
|
|
|
|
return result
|
|
|
|
|
|
# Url format is:
|
|
# s3(s)://<username>:<password>@<region>+<host:port>/<bucket name>/<path>
|
|
# WARNING: this regex is not RFC173-compliant, some parts (e.g. username)
|
|
# should be escaped but here are not.
|
|
S3_URL = re.compile(
|
|
# Start of the string
|
|
r"\A"
|
|
# s3:// (for HTTP) or s3s (for HTTPS)
|
|
r"^(?P<proto>s3(|s))://"
|
|
# username, followed by a ':'
|
|
r"(?P<username>[^:]*):"
|
|
# password, followed by a '@'
|
|
r"(?P<password>[^@]*)@"
|
|
# region, followed by a '+'
|
|
r"(?P<region>[^+]*)\+"
|
|
# host (including ':port'), followed by a '/'
|
|
r"(?P<host>[^/]*)/"
|
|
# bucket, which is the first path component and is required
|
|
r"(?P<bucket>[^/]+)"
|
|
# path, which includes a leading slash
|
|
r"(?P<path>/.*)?"
|
|
# End of the string
|
|
r"\Z"
|
|
)
|
|
|
|
|
|
class S3Synchronizer(Synchronizer):
|
|
def __init__(self, url: str):
|
|
self.credentials: str | None = None
|
|
match_obj = S3_URL.match(url)
|
|
if match_obj is None:
|
|
raise ValueError("S3 endpoint invalid")
|
|
|
|
self.botocore_session = BotocoreSession()
|
|
self.botocore_session.set_credentials(match_obj["username"], match_obj["password"])
|
|
config = botocore.config.Config(s3={"addressing_style": "path"})
|
|
session = boto3.session.Session(botocore_session=self.botocore_session)
|
|
proto = "https" if match_obj["proto"] == "s3s" else "http"
|
|
self.client = session.client(
|
|
"s3",
|
|
endpoint_url=f"{proto}://{match_obj['host']}",
|
|
region_name=match_obj["region"],
|
|
config=config,
|
|
)
|
|
self.bucket = match_obj["bucket"]
|
|
if match_obj["path"] is not None:
|
|
self.path = match_obj["path"][1:]
|
|
else:
|
|
self.path = ""
|
|
|
|
self.lock = Lock()
|
|
# The following properties are mutated when calling load or save,
|
|
# because of that they require the above lock to be acquired.
|
|
self.last_index: int = 0
|
|
self.last_save_event: Event | None = None
|
|
self.last_file_list: dict[str, FileData] = {}
|
|
|
|
def load(self, path: Path) -> bool:
|
|
"""Download the contents of S3 to the specified path. Since this method
|
|
is typically called at startup it is not thread-safe."""
|
|
|
|
# Check that the provided path is not a file
|
|
assert not path.exists() or path.is_dir()
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Download index.yml from S3 and parse it
|
|
index = self._get_index()
|
|
# If the index is empty do nothing
|
|
if len(index) == 0:
|
|
return True
|
|
|
|
# Generate the `last_file_list`, with the data retrieved from index.yml
|
|
self.last_file_list = {k: FileData(v) for k, v in index.items()}
|
|
# Download all the files mentioned in index.yml in the specified
|
|
# resume directory
|
|
queue = [
|
|
executor.submit(self._download_file, value, path / key) for key, value in index.items()
|
|
]
|
|
# Actually wait for all the downloads to finish
|
|
return gather_bools(queue, "Failed to load")
|
|
|
|
def save(self, path: Path) -> bool:
|
|
# Check that the provided path is not a file
|
|
assert not path.exists() or path.is_dir()
|
|
|
|
# Early exit, if the directory does not exist or file-list isn't there
|
|
if not path.exists() or not (path / "file-list.yml").is_file():
|
|
return True
|
|
|
|
# Read and parse the file list
|
|
with open(path / "file-list.yml") as f:
|
|
file_list = yaml.safe_load(f)
|
|
|
|
with self.lock:
|
|
# Make an early exit if the index is "behind". That means that
|
|
# either:
|
|
# * We've been called twice
|
|
# * We've been beaten to the punch by another thread
|
|
if file_list["Index"] <= self.last_index:
|
|
return True
|
|
|
|
# Acquire a position in the queue of committing `index.yml`:
|
|
# 1. Steal `last_save_event`: we will wait on it before committing
|
|
# our `index.yml`
|
|
# 2. Create a new event on which the next save will wait on (before
|
|
# committing their `index.yml`)
|
|
# 3. Push all the required files to S3, fail in case an error
|
|
# happens
|
|
# 4. Wait on the stolen `last_save_event`
|
|
# 5. Push our `index.yml`
|
|
# 6. Notify that we are done and the next save can proceed
|
|
# committing a new `index.yml`. This is done automatically by
|
|
# the `with_event` context manager
|
|
last_save_event = self.last_save_event
|
|
finish_event = Event()
|
|
self.last_save_event = finish_event
|
|
|
|
# Also save the previous_last_index, this accounts for situations
|
|
# where an index was skipped
|
|
previous_last_index = self.last_index
|
|
self.last_index = file_list["Index"]
|
|
|
|
with with_event(finish_event):
|
|
to_upload = set()
|
|
new_file_list = {}
|
|
for file_path, attributes in file_list["Files"].items():
|
|
# The if below decides if a file needs to be uploaded, which
|
|
# happens in one of two conditions:
|
|
# * The path was not seen in the last upload
|
|
# * The path has an index higher than last time
|
|
if attributes["Index"] > previous_last_index:
|
|
# Create a new filename string, this prepends a UUID to the
|
|
# filename, which preserving the path.
|
|
new_filename = self._generate_filename(file_path)
|
|
to_upload.add(file_path)
|
|
# Generate a new FileData entry with the index and encoding
|
|
# from file-list.yml
|
|
new_file_list[file_path] = FileData(
|
|
new_filename, attributes["Index"], attributes["Encoding"]
|
|
)
|
|
|
|
queue = []
|
|
for key in to_upload:
|
|
data = new_file_list[key]
|
|
queue.append(
|
|
executor.submit(self._upload_file, path / key, data.s3_filename, data.encoding)
|
|
)
|
|
|
|
if not gather_bools(queue, "Failed to save"):
|
|
return False
|
|
|
|
# Before writing the index file proper, wait for the previous event
|
|
# to finish
|
|
if last_save_event is not None:
|
|
last_save_event.wait()
|
|
|
|
with self.lock:
|
|
# Grab the entries in `file_list` which were saved in previous
|
|
# iterations.
|
|
for file_path, attributes in file_list["Files"].items():
|
|
if file_path not in new_file_list:
|
|
new_file_list[file_path] = self.last_file_list[file_path]
|
|
|
|
# Generate the `index.yml` file and upload it to S3
|
|
new_index = {k: v.s3_filename for k, v in new_file_list.items()}
|
|
with NamedTemporaryFile("w") as temp_index_file:
|
|
yaml.safe_dump(new_index, temp_index_file)
|
|
result = self._upload_file(Path(temp_index_file.name), "index.yml", "None")
|
|
if not result:
|
|
return False
|
|
|
|
# Replace last_file_list with the new one
|
|
self.last_file_list = new_file_list
|
|
|
|
return True
|
|
|
|
@staticmethod
|
|
def _generate_filename(filename: str) -> str:
|
|
filename_path = Path(filename)
|
|
return str(filename_path.parent / f"{uuid4()}-{filename_path.name}")
|
|
|
|
def set_credentials(self, credentials: str):
|
|
access_key, secret_key = credentials.split(":", 1)
|
|
self.botocore_session.set_credentials(access_key, secret_key)
|
|
|
|
def _get_key(self, name: str) -> str:
|
|
return os.path.join(self.path, name) if name != "" else self.path
|
|
|
|
def _download_file(self, key: str, destination: Path) -> bool:
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
key = self._get_key(key)
|
|
try:
|
|
self.client.download_file(
|
|
Key=key,
|
|
Bucket=self.bucket,
|
|
Filename=str(destination.resolve()),
|
|
)
|
|
return True
|
|
except self.client.exceptions.NoSuchKey:
|
|
return False
|
|
|
|
def _read_file(self, name: str) -> bytes | None:
|
|
key = self._get_key(name)
|
|
try:
|
|
object_request = self.client.get_object(Bucket=self.bucket, Key=key)
|
|
except self.client.exceptions.NoSuchKey:
|
|
return None
|
|
|
|
return cast(bytes, object_request["Body"].read())
|
|
|
|
def _get_index(self) -> dict[str, str]:
|
|
data = self._read_file("index.yml")
|
|
if data is None:
|
|
return {}
|
|
return yaml.safe_load(data.decode("utf-8"))
|
|
|
|
def _upload_file(self, source: Path, name: str, encoding: str) -> bool:
|
|
key = self._get_key(name)
|
|
extra_args = {"ACL": "private"}
|
|
if encoding == "Gzip":
|
|
extra_args["ContentEncoding"] = "gzip"
|
|
try:
|
|
self.client.upload_file(
|
|
Filename=str(source.resolve()),
|
|
Key=key,
|
|
Bucket=self.bucket,
|
|
ExtraArgs=extra_args,
|
|
)
|
|
return True
|
|
except boto3.exceptions.S3UploadFailedError as e:
|
|
sys.stderr.write("Failed uploading file\n")
|
|
traceback.print_exception(e, file=sys.stderr)
|
|
sys.stderr.flush()
|
|
return False
|