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

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

17ActionCacheService 

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

19 

20Allows clients to manually query/update the action cache. 

21""" 

22 

23import logging 

24from typing import Union, cast 

25 

26import grpc 

27 

28import buildgrid.server.context as context_module 

29from buildgrid._exceptions import NotFoundError 

30from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import DESCRIPTOR as RE_DESCRIPTOR 

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

32 ActionResult, 

33 GetActionResultRequest, 

34 UpdateActionResultRequest, 

35) 

36from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2_grpc import ( 

37 ActionCacheServicer, 

38 add_ActionCacheServicer_to_server, 

39) 

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

41from buildgrid.server.actioncache.instance import ActionCache 

42from buildgrid.server.auth.manager import authorize_unary_unary 

43from buildgrid.server.metrics_names import ( 

44 AC_CACHE_HITS_METRIC_NAME, 

45 AC_CACHE_MISSES_METRIC_NAME, 

46 AC_GET_ACTION_RESULT_TIME_METRIC_NAME, 

47 AC_UPDATE_ACTION_RESULT_TIME_METRIC_NAME, 

48) 

49from buildgrid.server.metrics_utils import DurationMetric, publish_counter_metric 

50from buildgrid.server.request_metadata_utils import printable_request_metadata 

51from buildgrid.server.servicer import InstancedServicer 

52from buildgrid.server.utils.decorators import handle_errors_unary_unary, track_request_id 

53 

54LOGGER = logging.getLogger(__name__) 

55 

56 

57class ActionCacheService(ActionCacheServicer, InstancedServicer[Union[ActionCache, ActionCacheABC]]): 

58 REGISTER_METHOD = add_ActionCacheServicer_to_server 

59 FULL_NAME = RE_DESCRIPTOR.services_by_name["ActionCache"].full_name 

60 

61 @authorize_unary_unary(lambda r: cast(str, r.instance_name)) 

62 @context_module.metadatacontext() 

63 @track_request_id 

64 @handle_errors_unary_unary(ActionResult) 

65 def GetActionResult(self, request: GetActionResultRequest, context: grpc.ServicerContext) -> ActionResult: 

66 LOGGER.info( 

67 f"GetActionResult request from [{context.peer()}] " 

68 f"([{printable_request_metadata(context.invocation_metadata())}])" 

69 ) 

70 

71 try: 

72 instance = self.get_instance(request.instance_name) 

73 with DurationMetric(AC_GET_ACTION_RESULT_TIME_METRIC_NAME, request.instance_name, instanced=True): 

74 action_result = instance.get_action_result(request.action_digest) 

75 publish_counter_metric(AC_CACHE_HITS_METRIC_NAME, 1, {"instance-name": request.instance_name}) 

76 return action_result 

77 

78 except NotFoundError: 

79 publish_counter_metric(AC_CACHE_MISSES_METRIC_NAME, 1, {"instance-name": request.instance_name}) 

80 raise 

81 

82 @authorize_unary_unary(lambda r: cast(str, r.instance_name)) 

83 @context_module.metadatacontext() 

84 @track_request_id 

85 @handle_errors_unary_unary(ActionResult) 

86 def UpdateActionResult(self, request: UpdateActionResultRequest, context: grpc.ServicerContext) -> ActionResult: 

87 LOGGER.info( 

88 f"UpdateActionResult request from [{context.peer()}] " 

89 f"([{printable_request_metadata(context.invocation_metadata())}])" 

90 ) 

91 

92 instance = self.get_instance(request.instance_name) 

93 with DurationMetric(AC_UPDATE_ACTION_RESULT_TIME_METRIC_NAME, request.instance_name, instanced=True): 

94 instance.update_action_result(request.action_digest, request.action_result) 

95 return request.action_result