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-06-11 15:37 +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. 

14 

15from asyncio import Lock 

16from collections import OrderedDict 

17from datetime import datetime, timedelta 

18from typing import Generic, Optional, TypeVar 

19 

20from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import ActionResult 

21from buildgrid._protos.google.longrunning.operations_pb2 import Operation 

22 

23T = TypeVar("T", ActionResult, Operation, bytes) 

24 

25 

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 

31 

32 def __str__(self) -> str: 

33 return str(self.value) 

34 

35 def __repr__(self) -> str: 

36 return repr(self.value) 

37 

38 def is_fresh(self) -> bool: 

39 return self._never_expire or datetime.utcnow() < self._expiry_date 

40 

41 

42class LruCache(Generic[T]): 

43 """Class implementing an async-safe LRU cache with an asynchronous API. 

44 

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. 

49 

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. 

53 

54 """ 

55 

56 def __init__(self, max_length: int): 

57 """Initialize a new LruCache. 

58 

59 Args: 

60 max_length (int): The maximum number of entries to store in the 

61 cache at any time. 

62 

63 """ 

64 self._cache: "OrderedDict[str, _CacheEntry[T]]" = OrderedDict() 

65 self._lock = Lock() 

66 self._max_length = max_length 

67 

68 def max_size(self) -> int: 

69 """Get the maximum number of items that can be stored in the cache. 

70 

71 Calling ``LruCache.set`` when there are already this many entries 

72 in the cache will cause the oldest entry to be dropped. 

73 

74 Returns: 

75 int: The maximum number of items to store. 

76 

77 """ 

78 return self._max_length 

79 

80 async def size(self) -> int: 

81 """Get the current number of items in the cache. 

82 

83 Returns: 

84 int: The number of items currently stored. 

85 

86 """ 

87 async with self._lock: 

88 return len(self._cache) 

89 

90 async def get(self, key: str) -> Optional[T]: 

91 """Get the value for a given key, or ``None``. 

92 

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. 

97 

98 Args: 

99 key (str): The key to get the corresponding value for. 

100 

101 Returns: 

102 A value mapped to the given key, or ``None`` 

103 

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 

114 

115 async def set(self, key: str, value: T, ttl: int) -> None: 

116 """Set the value and TTL of a key. 

117 

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. 

122 

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. 

127 

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)