Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/actioncache/caches/lru_cache.py: 100.00%

1# Copyright (C) 2020 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 

15 

16import collections 

17 

18from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2 

19from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import ActionResult, Digest 

20from buildgrid.server.actioncache.caches.action_cache_abc import ActionCacheABC 

21from buildgrid.server.cas.storage.storage_abc import StorageABC 

22from buildgrid.server.exceptions import NotFoundError 

23from buildgrid.server.logging import buildgrid_logger 

24 

25LOGGER = buildgrid_logger(__name__) 

26 

27 

28class LruActionCache(ActionCacheABC): 

29 """In-memory Action Cache implementation with LRU eviction. 

30 

31 This cache has a configurable fixed size, evicting the least recently 

32 accessed entry when adding a new entry would exceed the fixed size. The 

33 cache is entirely stored in memory so its contents are lost on restart. 

34 

35 This type of cache is ideal for use cases that need a simple and fast 

36 cache, with no requirements for longevity of the cache content. It is not 

37 recommended to use this type of cache in situations where you may wish to 

38 obtain cached results a reasonable time in the future, due to its fixed 

39 size. 

40 

41 """ 

42 

43 def __init__( 

44 self, 

45 storage: StorageABC, 

46 max_cached_refs: int, 

47 allow_updates: bool = True, 

48 cache_failed_actions: bool = True, 

49 ): 

50 """Initialise a new in-memory LRU Action Cache. 

51 

52 Args: 

53 storage (StorageABC): Storage backend instance to be used. 

54 max_cached_refs (int): Maximum number of entries to store in the cache. 

55 allow_updates (bool): Whether to allow writing to the cache. If false, 

56 this is a read-only cache for all clients. 

57 cache_failed_actions (bool): Whether or not to cache Actions with 

58 non-zero exit codes. 

59 

60 """ 

61 super().__init__(storage=storage) 

62 

63 self._cache_failed_actions = cache_failed_actions 

64 self._storage = storage 

65 self._allow_updates = allow_updates 

66 self._max_cached_refs = max_cached_refs 

67 self._digest_map = collections.OrderedDict() # type: ignore 

68 

69 def get_action_result(self, action_digest: Digest) -> ActionResult: 

70 """Retrieves the cached result for an Action. 

71 

72 If there is no cached result found, raises ``NotFoundError``. 

73 

74 Args: 

75 action_digest (Digest): The digest of the Action to retrieve the 

76 cached result of. 

77 

78 """ 

79 key = self._get_key(action_digest) 

80 if key in self._digest_map: 

81 assert self._storage, "Storage used before initialization" 

82 action_result = self._storage.get_message(self._digest_map[key], remote_execution_pb2.ActionResult) 

83 

84 if action_result is not None: 

85 if self.referenced_blobs_still_exist(action_digest, action_result): 

86 self._digest_map.move_to_end(key) 

87 return action_result 

88 

89 if self._allow_updates: 

90 LOGGER.debug( 

91 "Removing action digest from cache due to missing blobs in CAS.", 

92 tags=dict(digest=action_digest), 

93 ) 

94 del self._digest_map[key] 

95 

96 raise NotFoundError(f"Key not found: {key}") 

97 

98 def update_action_result(self, action_digest: Digest, action_result: ActionResult) -> None: 

99 """Stores a result for an Action in the cache. 

100 

101 If the result has a non-zero exit code and `cache_failed_actions` is False 

102 for this cache, the result is not cached. 

103 

104 Args: 

105 action_digest (Digest): The digest of the Action whose result is 

106 being cached. 

107 action_result (ActionResult): The result to cache for the given 

108 Action digest. 

109 

110 """ 

111 if self._cache_failed_actions or action_result.exit_code == 0: 

112 key = self._get_key(action_digest) 

113 if not self._allow_updates: 

114 raise NotImplementedError("Updating cache not allowed") 

115 

116 if self._max_cached_refs == 0: 

117 return 

118 

119 while len(self._digest_map) >= self._max_cached_refs: 

120 self._digest_map.popitem(last=False) 

121 

122 assert self._storage, "Storage used before initialization" 

123 result_digest = self._storage.put_message(action_result) 

124 self._digest_map[key] = result_digest 

125 

126 LOGGER.info("Result cached for action.", tags=dict(digest=action_digest)) 

127 

128 def _get_key(self, action_digest: Digest) -> tuple[str, int]: 

129 """Get a hashable cache key from a given Action digest. 

130 

131 Args: 

132 action_digest (Digest): The digest to produce a cache key for. 

133 

134 """ 

135 return (action_digest.hash, action_digest.size_bytes)