Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/actioncache/caches/remote_cache.py: 94.23%

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

16""" 

17Remote Action Cache 

18=================== 

19 

20Provides an interface to a remote Action Cache. This can be used by other 

21services (e.g. an Execution service) to communicate with a remote cache. 

22 

23It provides the same API as any other Action Cache instance backend. 

24 

25""" 

26 

27 

28import logging 

29from typing import Dict, Optional, Tuple 

30 

31from buildgrid._exceptions import GrpcUninitializedError, NotFoundError 

32from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import ( 

33 ActionResult, 

34 Digest 

35) 

36from buildgrid.client.actioncache import ActionCacheClient 

37from buildgrid.client.capabilities import CapabilitiesInterface 

38from buildgrid.client.channel import setup_channel 

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

40 

41 

42class RemoteActionCache(ActionCacheABC): 

43 

44 def __init__( 

45 self, 

46 remote: str, 

47 instance_name: str, 

48 channel_options: Optional[Tuple[Tuple[str, str], ...]]=None, 

49 tls_credentials: Optional[Dict[str, str]]=None 

50 ): 

51 """Initialises a new RemoteActionCache instance. 

52 

53 Args: 

54 remote (str): URL of the remote ActionCache service to open a 

55 channel to. 

56 instance_name (str): The instance name of the remote ActionCache 

57 service. 

58 channel_options (tuple): Optional tuple of channel options to set 

59 when opening the gRPC channel to the remote. 

60 tls_credentials (dict): Optional credentials to use when opening 

61 the gRPC channel. If unset then an insecure channel will be 

62 used. 

63 

64 """ 

65 ActionCacheABC.__init__(self) 

66 self._logger = logging.getLogger(__name__) 

67 self._remote_instance_name = instance_name 

68 self._remote = remote 

69 self._channel_options = channel_options 

70 if tls_credentials is None: 

71 tls_credentials = {} 

72 self._tls_credentials = tls_credentials 

73 self._channel = None 

74 self._allow_updates = None 

75 

76 self._action_cache = None 

77 

78 def setup_grpc(self): 

79 if self._channel is None: 

80 self._channel, _ = setup_channel( 

81 self._remote, 

82 auth_token=None, 

83 client_key=self._tls_credentials.get("tls-client-key"), 

84 client_cert=self._tls_credentials.get("tls-client-cert"), 

85 server_cert=self._tls_credentials.get("tls-server-cert") 

86 ) 

87 if self._action_cache is None: 

88 self._action_cache = ActionCacheClient( 

89 self._channel, self._remote_instance_name) 

90 

91 @property 

92 def allow_updates(self) -> bool: 

93 if self._channel is None: 

94 raise GrpcUninitializedError("Remote cache used before gRPC initialization.") 

95 

96 # Check if updates are allowed if we haven't already. 

97 # This is done the first time update_action_result is called rather 

98 # than on instantiation because the remote cache may not be running 

99 # when this object is instantiated. 

100 if self._allow_updates is None: 

101 interface = CapabilitiesInterface(self._channel) 

102 capabilities = interface.get_capabilities(self._remote_instance_name) 

103 self._allow_updates = (capabilities 

104 .cache_capabilities 

105 .action_cache_update_capabilities 

106 .update_enabled) 

107 return self._allow_updates 

108 

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

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

111 

112 Queries the remote ActionCache service to retrieve the cached 

113 result for a given Action digest. If the remote cache doesn't 

114 contain a result for the Action, then ``None`` is returned. 

115 

116 Args: 

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

118 cached result of. 

119 

120 """ 

121 if self._action_cache is None: 

122 raise GrpcUninitializedError("Remote cache used before gRPC initialization.") 

123 

124 action_result = self._action_cache.get(action_digest) 

125 

126 if action_result is None: 

127 key = self._get_key(action_digest) 

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

129 return action_result 

130 

131 def update_action_result(self, action_digest: Digest, 

132 action_result: ActionResult) -> None: 

133 """Stores a result for an Action in the remote cache. 

134 

135 Sends an ``UpdateActionResult`` request to the remote ActionCache 

136 service, to store the result in the remote cache. 

137 

138 If the remote cache doesn't allow updates, then this raises a 

139 ``NotImplementedError``. 

140 

141 Args: 

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

143 being cached. 

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

145 Action digest. 

146 

147 """ 

148 if self._action_cache is None: 

149 raise GrpcUninitializedError("Remote cache used before gRPC initialization.") 

150 

151 if not self.allow_updates: 

152 raise NotImplementedError("Updating cache not allowed") 

153 return self._action_cache.update(action_digest, action_result) 

154 

155 def _get_key(self, action_digest): 

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

157 

158 Args: 

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

160 

161 """ 

162 return (action_digest.hash, action_digest.size_bytes)