Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/utils/async_lru_cache.py: 92.86%
42 statements
« prev ^ index » next coverage.py v7.4.1, created at 2024-11-01 16:30 +0000
« prev ^ index » next coverage.py v7.4.1, created at 2024-11-01 16:30 +0000
1# Copyright (C) 2022 Bloomberg LP
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# <http://www.apache.org/licenses/LICENSE-2.0>
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15from asyncio import Lock
16from collections import OrderedDict
17from datetime import datetime, timedelta
18from typing import Generic, Optional, TypeVar
20from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import ActionResult
21from buildgrid._protos.google.longrunning.operations_pb2 import Operation
23T = TypeVar("T", ActionResult, Operation, bytes)
26class _CacheEntry(Generic[T]):
27 def __init__(self, value: T, ttl: int) -> None:
28 self._never_expire = ttl <= 0
29 self._expiry_date = datetime.utcnow() + timedelta(seconds=ttl)
30 self.value: T = value
32 def __str__(self) -> str:
33 return str(self.value)
35 def __repr__(self) -> str:
36 return repr(self.value)
38 def is_fresh(self) -> bool:
39 return self._never_expire or datetime.utcnow() < self._expiry_date
42class LruCache(Generic[T]):
43 """Class implementing an async-safe LRU cache with an asynchronous API.
45 This class provides LRU functionality similar to the existing LruCache
46 class, but providing an asynchronous API and using asynchronous locks
47 internally. This allows it to be used in asyncio coroutines without
48 blocking the event loop.
50 This class also supports setting a TTL on cache entries. Cleanup of
51 entries which have outlived their TTL is done lazily when ``LruCache.get``
52 is called for the relevant key.
54 """
56 def __init__(self, max_length: int):
57 """Initialize a new LruCache.
59 Args:
60 max_length (int): The maximum number of entries to store in the
61 cache at any time.
63 """
64 self._cache: "OrderedDict[str, _CacheEntry[T]]" = OrderedDict()
65 self._lock = Lock()
66 self._max_length = max_length
68 def max_size(self) -> int:
69 """Get the maximum number of items that can be stored in the cache.
71 Calling ``LruCache.set`` when there are already this many entries
72 in the cache will cause the oldest entry to be dropped.
74 Returns:
75 int: The maximum number of items to store.
77 """
78 return self._max_length
80 async def size(self) -> int:
81 """Get the current number of items in the cache.
83 Returns:
84 int: The number of items currently stored.
86 """
87 async with self._lock:
88 return len(self._cache)
90 async def get(self, key: str) -> Optional[T]:
91 """Get the value for a given key, or ``None``.
93 This method returns the value for a given key. If the key isn't
94 in the cache, or the value's TTL has expired, then this returns
95 ``None`` instead. In the case of a TTL expiry, the key is removed
96 from the cache to save space.
98 Args:
99 key (str): The key to get the corresponding value for.
101 Returns:
102 A value mapped to the given key, or ``None``
104 """
105 async with self._lock:
106 entry = self._cache.get(key)
107 if entry is not None:
108 if entry.is_fresh():
109 self._cache.move_to_end(key)
110 return entry.value
111 else:
112 del self._cache[key]
113 return None
115 async def set(self, key: str, value: T, ttl: int) -> None:
116 """Set the value and TTL of a key.
118 This method sets the value and TTL of a key. A TTL of 0
119 or lower means that the key won't expire due to age. Keys
120 with a TTL of 0 or lower will still be dropped when they're
121 the least-recently-used key.
123 Args:
124 key (str): The key to update.
125 value (T): The value to map to the key.
126 ttl (int): A TTL (in seconds) for the key.
128 """
129 async with self._lock:
130 while len(self._cache) >= self._max_length:
131 self._cache.popitem(last=False)
132 self._cache[key] = _CacheEntry(value, ttl)