buffer.py
# Copyright (C) 2019 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information
from collections import deque
from functools import partial
from typing import Optional, Iterable, Dict
from swh.core.utils import grouper
from swh.model.model import Content, BaseModel
from swh.storage import get_storage
class BufferingProxyStorage:
"""Storage implementation in charge of accumulating objects prior to
discussing with the "main" storage.
Sample configuration use case for buffering storage:
.. code-block:: yaml
storage:
cls: buffer
args:
storage:
cls: remote
args: http://storage.internal.staging.swh.network:5002/
min_batch_size:
content: 10000
content_bytes: 100000000
skipped_content: 10000
directory: 5000
revision: 1000
release: 10000
"""
def __init__(self, storage, min_batch_size=None):
self.storage = get_storage(**storage)
if min_batch_size is None:
min_batch_size = {}
self.min_batch_size = {
'content': min_batch_size.get('content', 10000),
'content_bytes': min_batch_size.get('content_bytes',
100*1024*1024),
'skipped_content': min_batch_size.get('skipped_content', 10000),
'directory': min_batch_size.get('directory', 25000),
'revision': min_batch_size.get('revision', 100000),
'release': min_batch_size.get('release', 100000),
}
self.object_types = [
'content', 'skipped_content', 'directory', 'revision', 'release']
self._objects = {k: deque() for k in self.object_types}
def __getattr__(self, key):
if key.endswith('_add'):
object_type = key.rsplit('_', 1)[0]
if object_type in self.object_types:
return partial(
self.object_add, object_type=object_type
)
if key == 'storage':
raise AttributeError(key)
return getattr(self.storage, key)
def content_add(self, content: Iterable[Content]) -> Dict:
"""Enqueue contents to write to the storage.
Following policies apply:
- First, check if the queue's threshold is hit.
If it is flush content to the storage.
- If not, check if the total size of enqueued contents's
threshold is hit. If it is flush content to the storage.
"""
content = list(content)
s = self.object_add(content, object_type='content')
if not s:
q = self._objects['content']
total_size = sum(c.length for c in q)
if total_size >= self.min_batch_size['content_bytes']:
return self.flush(['content'])
return s
def flush(self, object_types: Optional[Iterable[str]] = None) -> Dict:
if object_types is None:
object_types = self.object_types
summary = {} # type: Dict[str, Dict]
for object_type in object_types:
q = self._objects[object_type]
for objs in grouper(q, n=self.min_batch_size[object_type]):
add_fn = getattr(self.storage, '%s_add' % object_type)
s = add_fn(objs)
summary = {k: v + summary.get(k, 0)
for k, v in s.items()}
q.clear()
return summary
def object_add(
self, objects: Iterable[BaseModel], *, object_type: str) -> Dict:
"""Enqueue objects to write to the storage. This checks if the queue's
threshold is hit. If it is actually write those to the storage.
"""
q = self._objects[object_type]
threshold = self.min_batch_size[object_type]
q.extend(objects)
if len(q) >= threshold:
return self.flush()
return {}