312 lines
9.3 KiB
Python
312 lines
9.3 KiB
Python
from __future__ import annotations
|
|
|
|
import abc
|
|
import dataclasses
|
|
import threading
|
|
from collections import OrderedDict, deque
|
|
from contextvars import ContextVar
|
|
from typing import (
|
|
Callable,
|
|
Deque,
|
|
Dict,
|
|
Optional,
|
|
Protocol,
|
|
Union,
|
|
)
|
|
|
|
from .core import Domain, PartialResult, Resolver
|
|
|
|
__all__ = [
|
|
"Cache",
|
|
"CachingResolver",
|
|
"Lru",
|
|
"S3Fifo",
|
|
"Sieve",
|
|
]
|
|
|
|
|
|
class Cache(Protocol):
|
|
"""Cache()
|
|
|
|
Cache abstract protocol. The :class:`CachingResolver` will look
|
|
values up, merge what was returned (possibly nothing) with what it
|
|
got from its actual parser, and *re-set the result*.
|
|
"""
|
|
|
|
@abc.abstractmethod
|
|
def __setitem__(self, key: str, value: PartialResult) -> None:
|
|
"""Adds or replace ``value`` to the cache at key ``key``."""
|
|
...
|
|
|
|
@abc.abstractmethod
|
|
def __getitem__(self, key: str) -> Optional[PartialResult]:
|
|
"""Returns a partial result for ``key`` if there is any."""
|
|
...
|
|
|
|
|
|
class Lru:
|
|
"""Cache following a least-recently used replacement policy: when
|
|
there is no more room in the cache, whichever entry was last seen
|
|
the least recently is removed.
|
|
|
|
Simple LRUs are generally outdated and to avoid as they have
|
|
relatively low hit rates for modern caches (at lower sizes). The
|
|
main use case here is if the workload can lead to the cache being
|
|
full of popular items then all of them being replaced at once:
|
|
:class:`S3Fifo` and :class:`Sieve` are FIFO-based caches and have
|
|
worst-case O(n) eviction.
|
|
"""
|
|
|
|
def __init__(self, maxsize: int):
|
|
self.maxsize = maxsize
|
|
self.cache: OrderedDict[str, PartialResult] = OrderedDict()
|
|
self.lock = threading.Lock()
|
|
|
|
def __getitem__(self, key: str) -> Optional[PartialResult]:
|
|
with self.lock:
|
|
e = self.cache.get(key)
|
|
if e:
|
|
self.cache.move_to_end(key)
|
|
return e
|
|
|
|
def __setitem__(self, key: str, value: PartialResult) -> None:
|
|
with self.lock:
|
|
if len(self.cache) >= self.maxsize and key not in self.cache:
|
|
self.cache.popitem(last=False)
|
|
self.cache[key] = value
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class CacheEntry:
|
|
__slots__ = ["freq", "key", "value"]
|
|
key: str
|
|
value: PartialResult
|
|
freq: int
|
|
|
|
|
|
class S3Fifo:
|
|
"""FIFO-based quick-demotion lazy-promotion cache [S3-FIFO]_.
|
|
|
|
Experimentally provides excellent hit rate at lower cache sizes,
|
|
for a relatively simple and efficient implementation. Notably
|
|
excellent at handling "one hit wonders", aka entries seen only
|
|
once during a work-set (or reasonable work window).
|
|
"""
|
|
|
|
def __init__(self, maxsize: int):
|
|
self.maxsize = maxsize
|
|
self.index: Dict[str, Union[CacheEntry, str]] = {}
|
|
self.small_target = max(1, int(maxsize / 10))
|
|
self.small: Deque[CacheEntry] = deque()
|
|
self.main_target = maxsize - self.small_target
|
|
self.main: Deque[CacheEntry] = deque()
|
|
self.ghost: Deque[str] = deque()
|
|
self.lock = threading.Lock()
|
|
|
|
def __getitem__(self, key: str) -> Optional[PartialResult]:
|
|
if (e := self.index.get(key)) and type(e) is CacheEntry:
|
|
# small race here, we could miss an increment
|
|
e.freq = min(e.freq + 1, 3)
|
|
return e.value
|
|
|
|
return None
|
|
|
|
def __setitem__(self, key: str, r: PartialResult) -> None:
|
|
with self.lock:
|
|
if (e := self.index.get(key)) and type(e) is CacheEntry:
|
|
e.value = r
|
|
return
|
|
|
|
if len(self.small) + len(self.main) >= self.maxsize:
|
|
# if main is not overcapacity, resize small
|
|
if len(self.main) < self.main_target:
|
|
self._evict_small()
|
|
# evict_small could have moved every entry to main, in
|
|
# which case we now need to evict from main
|
|
if len(self.small) + len(self.main) >= self.maxsize:
|
|
self._evict_main()
|
|
|
|
entry = CacheEntry(key, r, 0)
|
|
if type(self.index.get(key)) is str:
|
|
self.main.appendleft(entry)
|
|
else:
|
|
self.small.appendleft(entry)
|
|
self.index[key] = entry
|
|
|
|
def _evict_main(self) -> None:
|
|
while True:
|
|
e = self.main.pop()
|
|
if e.freq:
|
|
e.freq -= 1
|
|
self.main.appendleft(e)
|
|
else:
|
|
del self.index[e.key]
|
|
return
|
|
|
|
def _evict_small(self) -> None:
|
|
while self.small:
|
|
e = self.small.pop()
|
|
if e.freq:
|
|
e.freq = 0
|
|
self.main.appendleft(e)
|
|
else:
|
|
g = self.index[e.key] = e.key
|
|
self.ghost.appendleft(g)
|
|
while len(self.ghost) > self.main_target:
|
|
g = self.ghost.pop()
|
|
if self.index.get(g) is g:
|
|
del self.index[g]
|
|
return
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class SieveNode:
|
|
__slots__ = ("key", "next", "value", "visited")
|
|
key: str
|
|
value: PartialResult
|
|
visited: bool
|
|
next: Optional[SieveNode]
|
|
|
|
|
|
class Sieve:
|
|
"""FIFO-based quick-demotion cache [SIEVE]_.
|
|
|
|
Simpler FIFO-based cache, cousin of :class:`S3Fifo`.
|
|
Experimentally slightly lower hit rates than :class:`S3Fifo` (if
|
|
way superior to LRU still), but much more compact (~50% lower
|
|
memory overhead at larger cache sizes, up to 100% at very small
|
|
cache sizes).
|
|
|
|
Can be an interesting candidate when trying to save on memory,
|
|
although the contained entries will generally be much larger than
|
|
the cache itself.
|
|
"""
|
|
|
|
def __init__(self, maxsize: int) -> None:
|
|
self.maxsize = maxsize
|
|
self.cache: Dict[str, SieveNode] = {}
|
|
self.head: Optional[SieveNode] = None
|
|
self.tail: Optional[SieveNode] = None
|
|
self.hand: Optional[SieveNode] = None
|
|
self.prev: Optional[SieveNode] = None
|
|
self.lock = threading.Lock()
|
|
|
|
def __getitem__(self, key: str) -> Optional[PartialResult]:
|
|
if entry := self.cache.get(key):
|
|
entry.visited = True
|
|
return entry.value
|
|
|
|
return None
|
|
|
|
def __setitem__(self, key: str, value: PartialResult) -> None:
|
|
with self.lock:
|
|
if e := self.cache.get(key):
|
|
e.value = value
|
|
return
|
|
|
|
if len(self.cache) >= self.maxsize:
|
|
self._evict()
|
|
|
|
node = self.cache[key] = SieveNode(key, value, False, None)
|
|
if self.head:
|
|
self.head.next = node
|
|
self.head = node
|
|
if self.tail is None:
|
|
self.tail = node
|
|
|
|
def _evict(self) -> None:
|
|
obj: Optional[SieveNode]
|
|
if self.hand:
|
|
obj, pobj = self.hand, self.prev
|
|
else:
|
|
obj, pobj = self.tail, None
|
|
|
|
while obj and obj.visited:
|
|
obj.visited = False
|
|
if obj.next:
|
|
obj, pobj = obj.next, obj
|
|
else:
|
|
obj, pobj = self.tail, None
|
|
|
|
if not obj:
|
|
return
|
|
|
|
self.hand = obj.next
|
|
self.prev = pobj
|
|
|
|
del self.cache[obj.key]
|
|
if not obj.next:
|
|
self.head = pobj
|
|
|
|
if pobj:
|
|
pobj.next = obj.next
|
|
else:
|
|
self.tail = obj.next
|
|
|
|
|
|
class Local:
|
|
"""Thread local cache decorator. Takes a cache factory and lazily
|
|
instantiates a cache for each thread it's accessed from.
|
|
|
|
This means the cache capacity and memory consumption is
|
|
figuratively multiplied by however many threads the cache is used
|
|
from, but those threads don't share their caching, and thus don't
|
|
contend on cache use.
|
|
|
|
"""
|
|
|
|
def __init__(self, factory: Callable[[], Cache]) -> None:
|
|
self.cv: ContextVar[Cache] = ContextVar("local-cache")
|
|
self.factory = factory
|
|
|
|
@property
|
|
def cache(self) -> Cache:
|
|
c = self.cv.get(None)
|
|
if c is None:
|
|
c = self.factory()
|
|
self.cv.set(c)
|
|
return c
|
|
|
|
def __getitem__(self, key: str) -> Optional[PartialResult]:
|
|
return self.cache[key]
|
|
|
|
def __setitem__(self, key: str, value: PartialResult) -> None:
|
|
self.cache[key] = value
|
|
|
|
|
|
class CachingResolver:
|
|
"""A wrapper resolver which takes an underlying concrete
|
|
:class:`Cache` for the actual caching and cache strategy.
|
|
|
|
This resolver only interacts with the :class:`Cache` and delegates
|
|
to the wrapped resolver in case of lookup failure.
|
|
|
|
:class:`CachingParser` will set entries back in the cache when
|
|
filling them up, it does not update results in place (and can't
|
|
really, they're immutable).
|
|
|
|
"""
|
|
|
|
def __init__(self, resolver: Resolver, cache: Cache):
|
|
self.parser: Resolver = resolver
|
|
self.cache: Cache = cache
|
|
|
|
def __call__(self, ua: str, domains: Domain, /) -> PartialResult:
|
|
entry = self.cache[ua]
|
|
if entry:
|
|
if domains in entry.domains:
|
|
return entry
|
|
|
|
domains &= ~entry.domains
|
|
|
|
r = self.parser(ua, domains)
|
|
if entry:
|
|
r = PartialResult(
|
|
string=ua,
|
|
domains=entry.domains | r.domains,
|
|
user_agent=entry.user_agent or r.user_agent,
|
|
os=entry.os or r.os,
|
|
device=entry.device or r.device,
|
|
)
|
|
self.cache[ua] = r
|
|
return r
|