wehub-resource-sync
116 lines
4.2 KiB
Python
116 lines
4.2 KiB
Python
import logging
|
|
from typing import Callable, Optional, Set
|
|
|
|
default_logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_MAX_URI_CACHE_SIZE_BYTES = (1024**3) * 10 # 10 GB
|
|
|
|
|
|
class URICache:
|
|
"""Caches URIs up to a specified total size limit.
|
|
|
|
URIs are represented by strings. Each URI has an associated size on disk.
|
|
When a URI is added to the URICache, it is marked as "in use".
|
|
When a URI is no longer in use, the user of this class should call
|
|
`mark_unused` to signal that the URI is safe for deletion.
|
|
|
|
URIs in the cache can be marked as "in use" by calling `mark_used`.
|
|
|
|
Deletion of URIs on disk does not occur until the size limit is exceeded.
|
|
When this happens, URIs that are not in use are deleted randomly until the
|
|
size limit is satisfied, or there are no more URIs that are not in use.
|
|
|
|
It is possible for the total size on disk to exceed the size limit if all
|
|
the URIs are in use.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
delete_fn: Optional[Callable[[str, logging.Logger], int]] = None,
|
|
max_total_size_bytes: int = DEFAULT_MAX_URI_CACHE_SIZE_BYTES,
|
|
debug_mode: bool = False,
|
|
):
|
|
# Maps URIs to the size in bytes of their corresponding disk contents.
|
|
self._used_uris: Set[str] = set()
|
|
self._unused_uris: Set[str] = set()
|
|
|
|
if delete_fn is None:
|
|
self._delete_fn = lambda uri, logger: 0
|
|
else:
|
|
self._delete_fn = delete_fn
|
|
|
|
# Total size of both used and unused URIs in the cache.
|
|
self._total_size_bytes = 0
|
|
self.max_total_size_bytes = max_total_size_bytes
|
|
|
|
# Used in `self._check_valid()` for testing.
|
|
self._debug_mode = debug_mode
|
|
|
|
def mark_unused(self, uri: str, logger: logging.Logger = default_logger):
|
|
"""Mark a URI as unused and okay to be deleted."""
|
|
if uri not in self._used_uris:
|
|
logger.info(f"URI {uri} is already unused.")
|
|
else:
|
|
self._unused_uris.add(uri)
|
|
self._used_uris.remove(uri)
|
|
logger.info(f"Marked URI {uri} unused.")
|
|
self._evict_if_needed(logger)
|
|
self._check_valid()
|
|
|
|
def mark_used(self, uri: str, logger: logging.Logger = default_logger):
|
|
"""Mark a URI as in use. URIs in use will not be deleted."""
|
|
if uri in self._used_uris:
|
|
return
|
|
elif uri in self._unused_uris:
|
|
self._used_uris.add(uri)
|
|
self._unused_uris.remove(uri)
|
|
else:
|
|
raise ValueError(
|
|
f"Got request to mark URI {uri} used, but this "
|
|
"URI is not present in the cache."
|
|
)
|
|
logger.info(f"Marked URI {uri} used.")
|
|
self._check_valid()
|
|
|
|
def add(self, uri: str, size_bytes: int, logger: logging.Logger = default_logger):
|
|
"""Add a URI to the cache and mark it as in use."""
|
|
if uri in self._unused_uris:
|
|
self._unused_uris.remove(uri)
|
|
|
|
self._used_uris.add(uri)
|
|
self._total_size_bytes += size_bytes
|
|
|
|
self._evict_if_needed(logger)
|
|
self._check_valid()
|
|
logger.info(f"Added URI {uri} with size {size_bytes}")
|
|
|
|
def get_total_size_bytes(self) -> int:
|
|
return self._total_size_bytes
|
|
|
|
def _evict_if_needed(self, logger: logging.Logger = default_logger):
|
|
"""Evict unused URIs (if they exist) until total size <= max size."""
|
|
while (
|
|
self._unused_uris
|
|
and self.get_total_size_bytes() > self.max_total_size_bytes
|
|
):
|
|
# TODO(architkulkarni): Evict least recently used URI instead
|
|
arbitrary_unused_uri = next(iter(self._unused_uris))
|
|
self._unused_uris.remove(arbitrary_unused_uri)
|
|
num_bytes_deleted = self._delete_fn(arbitrary_unused_uri, logger)
|
|
self._total_size_bytes -= num_bytes_deleted
|
|
logger.info(
|
|
f"Deleted URI {arbitrary_unused_uri} with size " f"{num_bytes_deleted}."
|
|
)
|
|
|
|
def _check_valid(self):
|
|
"""(Debug mode only) Check "used" and "unused" sets are disjoint."""
|
|
if self._debug_mode:
|
|
assert self._used_uris & self._unused_uris == set()
|
|
|
|
def __contains__(self, uri):
|
|
return uri in self._used_uris or uri in self._unused_uris
|
|
|
|
def __repr__(self):
|
|
return str(self.__dict__)
|