Changeset View
Changeset View
Standalone View
Standalone View
swh/storage/proxies/buffer.py
| # Copyright (C) 2019-2020 The Software Heritage developers | # Copyright (C) 2019-2020 The Software Heritage developers | ||||
| # See the AUTHORS file at the top-level directory of this distribution | # See the AUTHORS file at the top-level directory of this distribution | ||||
| # License: GNU General Public License version 3, or any later version | # License: GNU General Public License version 3, or any later version | ||||
| # See top-level LICENSE file for more information | # See top-level LICENSE file for more information | ||||
| from functools import partial | from functools import partial | ||||
| from typing import Dict, Iterable, Mapping, Sequence, Tuple | from typing import Dict, Iterable, Mapping, Sequence, Tuple | ||||
| from typing_extensions import Literal | from typing_extensions import Literal | ||||
| from swh.core.utils import grouper | from swh.core.utils import grouper | ||||
| from swh.model.model import BaseModel, Content, SkippedContent | from swh.model.model import BaseModel, Content, Directory, SkippedContent | ||||
| from swh.storage import get_storage | from swh.storage import get_storage | ||||
| from swh.storage.interface import StorageInterface | from swh.storage.interface import StorageInterface | ||||
| LObjectType = Literal[ | LObjectType = Literal[ | ||||
| "content", | "content", | ||||
| "skipped_content", | "skipped_content", | ||||
| "directory", | "directory", | ||||
| "revision", | "revision", | ||||
| Show All 11 Lines | OBJECT_TYPES: Tuple[LObjectType, ...] = ( | ||||
| "extid", | "extid", | ||||
| ) | ) | ||||
| DEFAULT_BUFFER_THRESHOLDS: Dict[str, int] = { | DEFAULT_BUFFER_THRESHOLDS: Dict[str, int] = { | ||||
| "content": 10000, | "content": 10000, | ||||
| "content_bytes": 100 * 1024 * 1024, | "content_bytes": 100 * 1024 * 1024, | ||||
| "skipped_content": 10000, | "skipped_content": 10000, | ||||
| "directory": 25000, | "directory": 25000, | ||||
| "directory_entries": 200000, | |||||
| "revision": 100000, | "revision": 100000, | ||||
| "release": 100000, | "release": 100000, | ||||
| "snapshot": 25000, | "snapshot": 25000, | ||||
| "extid": 10000, | "extid": 10000, | ||||
| } | } | ||||
| class BufferingProxyStorage: | class BufferingProxyStorage: | ||||
| Show All 12 Lines | .. code-block:: yaml | ||||
| storage: | storage: | ||||
| cls: remote | cls: remote | ||||
| args: http://storage.internal.staging.swh.network:5002/ | args: http://storage.internal.staging.swh.network:5002/ | ||||
| min_batch_size: | min_batch_size: | ||||
| content: 10000 | content: 10000 | ||||
| content_bytes: 100000000 | content_bytes: 100000000 | ||||
| skipped_content: 10000 | skipped_content: 10000 | ||||
| directory: 5000 | directory: 5000 | ||||
| directory_entries: 100000 | |||||
| revision: 1000 | revision: 1000 | ||||
| release: 10000 | release: 10000 | ||||
| snapshot: 5000 | snapshot: 5000 | ||||
| """ | """ | ||||
| def __init__(self, storage: Mapping, min_batch_size: Mapping = {}): | def __init__(self, storage: Mapping, min_batch_size: Mapping = {}): | ||||
| self.storage: StorageInterface = get_storage(**storage) | self.storage: StorageInterface = get_storage(**storage) | ||||
| self._buffer_thresholds = {**DEFAULT_BUFFER_THRESHOLDS, **min_batch_size} | self._buffer_thresholds = {**DEFAULT_BUFFER_THRESHOLDS, **min_batch_size} | ||||
| self._objects: Dict[LObjectType, Dict[Tuple[str, ...], BaseModel]] = { | self._objects: Dict[LObjectType, Dict[Tuple[str, ...], BaseModel]] = { | ||||
| k: {} for k in OBJECT_TYPES | k: {} for k in OBJECT_TYPES | ||||
| } | } | ||||
| self._contents_size: int = 0 | self._contents_size: int = 0 | ||||
| self._directory_entries: int = 0 | |||||
| def __getattr__(self, key: str): | def __getattr__(self, key: str): | ||||
| if key.endswith("_add"): | if key.endswith("_add"): | ||||
| object_type = key.rsplit("_", 1)[0] | object_type = key.rsplit("_", 1)[0] | ||||
| if object_type in OBJECT_TYPES: | if object_type in OBJECT_TYPES: | ||||
| return partial(self.object_add, object_type=object_type, keys=["id"],) | return partial(self.object_add, object_type=object_type, keys=["id"],) | ||||
| if key == "storage": | if key == "storage": | ||||
| raise AttributeError(key) | raise AttributeError(key) | ||||
| Show All 9 Lines | def content_add(self, contents: Sequence[Content]) -> Dict[str, int]: | ||||
| flush content to the storage. | flush content to the storage. | ||||
| """ | """ | ||||
| stats = self.object_add( | stats = self.object_add( | ||||
| contents, | contents, | ||||
| object_type="content", | object_type="content", | ||||
| keys=["sha1", "sha1_git", "sha256", "blake2s256"], | keys=["sha1", "sha1_git", "sha256", "blake2s256"], | ||||
| ) | ) | ||||
| if not stats: # We did not flush already | if not stats: | ||||
| # We did not flush based on number of objects; check total size | |||||
| self._contents_size += sum(c.length for c in contents) | self._contents_size += sum(c.length for c in contents) | ||||
| if self._contents_size >= self._buffer_thresholds["content_bytes"]: | if self._contents_size >= self._buffer_thresholds["content_bytes"]: | ||||
| return self.flush(["content"]) | return self.flush(["content"]) | ||||
| return stats | return stats | ||||
| def skipped_content_add(self, contents: Sequence[SkippedContent]) -> Dict[str, int]: | def skipped_content_add(self, contents: Sequence[SkippedContent]) -> Dict[str, int]: | ||||
| return self.object_add( | return self.object_add( | ||||
| contents, | contents, | ||||
| object_type="skipped_content", | object_type="skipped_content", | ||||
| keys=["sha1", "sha1_git", "sha256", "blake2s256"], | keys=["sha1", "sha1_git", "sha256", "blake2s256"], | ||||
| ) | ) | ||||
| def directory_add(self, directories: Sequence[Directory]) -> Dict[str, int]: | |||||
| stats = self.object_add(directories, object_type="directory", keys=["id"]) | |||||
| if not stats: | |||||
| # We did not flush based on number of objects; check the number of entries | |||||
| self._directory_entries += sum(len(d.entries) for d in directories) | |||||
| if self._directory_entries >= self._buffer_thresholds["directory_entries"]: | |||||
| return self.flush(["content", "directory"]) | |||||
| return stats | |||||
| def object_add( | def object_add( | ||||
| self, | self, | ||||
| objects: Sequence[BaseModel], | objects: Sequence[BaseModel], | ||||
| *, | *, | ||||
| object_type: LObjectType, | object_type: LObjectType, | ||||
| keys: Iterable[str], | keys: Iterable[str], | ||||
| ) -> Dict[str, int]: | ) -> Dict[str, int]: | ||||
| """Push objects to write to the storage in the buffer. Flushes the | """Push objects to write to the storage in the buffer. Flushes the | ||||
| ▲ Show 20 Lines • Show All 44 Lines • ▼ Show 20 Lines | def clear_buffers(self, object_types: Sequence[LObjectType] = OBJECT_TYPES) -> None: | ||||
| you want to continue your processing. | you want to continue your processing. | ||||
| """ | """ | ||||
| for object_type in object_types: | for object_type in object_types: | ||||
| buffer_ = self._objects[object_type] | buffer_ = self._objects[object_type] | ||||
| buffer_.clear() | buffer_.clear() | ||||
| if object_type == "content": | if object_type == "content": | ||||
| self._contents_size = 0 | self._contents_size = 0 | ||||
| elif object_type == "directory": | |||||
| self._directory_entries = 0 | |||||
| self.storage.clear_buffers(object_types) | self.storage.clear_buffers(object_types) | ||||