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 

17from typing import Tuple 

18 

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

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

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

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

23from buildgrid.server.exceptions import NotFoundError 

24from buildgrid.server.logging import buildgrid_logger 

25 

26LOGGER = buildgrid_logger(__name__) 

27 

28 

29class LruActionCache(ActionCacheABC): 

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

31 

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

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

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

35 

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

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

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

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

40 size. 

41 

42 """ 

43 

44 def __init__( 

45 self, storage: StorageABC, max_cached_refs: int, allow_updates: bool = True, cache_failed_actions: bool = True 

46 ): 

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

48 

49 Args: 

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

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

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

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

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

55 non-zero exit codes. 

56 

57 """ 

58 super().__init__(storage=storage) 

59 

60 self._cache_failed_actions = cache_failed_actions 

61 self._storage = storage 

62 self._allow_updates = allow_updates 

63 self._max_cached_refs = max_cached_refs 

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

65 

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

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

68 

69 If there is no cached result found, returns None. 

70 

71 Args: 

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

73 cached result of. 

74 

75 """ 

76 key = self._get_key(action_digest) 

77 if key in self._digest_map: 

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

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

80 

81 if action_result is not None: 

82 if self.referenced_blobs_still_exist(action_digest, action_result): 

83 self._digest_map.move_to_end(key) 

84 return action_result 

85 

86 if self._allow_updates: 

87 LOGGER.debug( 

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

89 tags=dict(digest=action_digest), 

90 ) 

91 del self._digest_map[key] 

92 

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

94 

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

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

97 

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

99 for this cache, the result is not cached. 

100 

101 Args: 

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

103 being cached. 

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

105 Action digest. 

106 

107 """ 

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

109 key = self._get_key(action_digest) 

110 if not self._allow_updates: 

111 raise NotImplementedError("Updating cache not allowed") 

112 

113 if self._max_cached_refs == 0: 

114 return 

115 

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

117 self._digest_map.popitem(last=False) 

118 

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

120 result_digest = self._storage.put_message(action_result) 

121 self._digest_map[key] = result_digest 

122 

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

124 

125 def _get_key(self, action_digest: Digest) -> Tuple[str, int]: 

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

127 

128 Args: 

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

130 

131 """ 

132 return (action_digest.hash, action_digest.size_bytes)