Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/bots/service.py: 95.65%

69 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-06-11 15:37 +0000

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""" 

17BotsService 

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

19 

20""" 

21 

22import logging 

23from typing import cast 

24 

25import grpc 

26from google.protobuf import empty_pb2 

27 

28from buildgrid._enums import BotStatus 

29from buildgrid._exceptions import InvalidArgumentError 

30from buildgrid._protos.google.devtools.remoteworkers.v1test2.bots_pb2 import DESCRIPTOR as BOTS_DESCRIPTOR 

31from buildgrid._protos.google.devtools.remoteworkers.v1test2.bots_pb2 import ( 

32 BotSession, 

33 CreateBotSessionRequest, 

34 PostBotEventTempRequest, 

35 UpdateBotSessionRequest, 

36) 

37from buildgrid._protos.google.devtools.remoteworkers.v1test2.bots_pb2_grpc import ( 

38 BotsServicer, 

39 add_BotsServicer_to_server, 

40) 

41from buildgrid.server.auth.manager import authorize 

42from buildgrid.server.bots.instance import BotsInterface 

43from buildgrid.server.instance import instanced 

44from buildgrid.server.metrics_names import ( 

45 BOTS_CREATE_BOT_SESSION_TIME_METRIC_NAME, 

46 BOTS_UPDATE_BOT_SESSION_TIME_METRIC_NAME, 

47) 

48from buildgrid.server.metrics_utils import DurationMetric 

49from buildgrid.server.servicer import InstancedServicer 

50from buildgrid.server.utils.context import CancellationContext 

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

52 

53LOGGER = logging.getLogger(__name__) 

54 

55 

56def _instance_name_from_bot_name(name: str) -> str: 

57 names = name.split("/") 

58 return "/".join(names[:-1]) 

59 

60 

61class BotsService(BotsServicer, InstancedServicer[BotsInterface]): 

62 REGISTER_METHOD = add_BotsServicer_to_server 

63 FULL_NAME = BOTS_DESCRIPTOR.services_by_name["Bots"].full_name 

64 

65 @instanced(lambda r: cast(str, r.parent)) 

66 @authorize 

67 @track_request_id 

68 @DurationMetric(BOTS_CREATE_BOT_SESSION_TIME_METRIC_NAME) 

69 @handle_errors_unary_unary(BotSession) 

70 def CreateBotSession(self, request: CreateBotSessionRequest, context: grpc.ServicerContext) -> BotSession: 

71 """Handles CreateBotSessionRequest messages. 

72 

73 Args: 

74 request (CreateBotSessionRequest): The incoming RPC request. 

75 context (grpc.ServicerContext): Context for the RPC call. 

76 """ 

77 LOGGER.info(f"CreateBotSession request from [{context.peer()}]") 

78 

79 instance_name = request.parent 

80 instance = self.get_instance(instance_name) 

81 bot_session = instance.create_bot_session( 

82 request.bot_session, CancellationContext(context), deadline=context.time_remaining() 

83 ) 

84 return bot_session 

85 

86 @instanced(lambda r: _instance_name_from_bot_name(r.name)) 

87 @authorize 

88 @track_request_id 

89 @DurationMetric(BOTS_UPDATE_BOT_SESSION_TIME_METRIC_NAME) 

90 @handle_errors_unary_unary(BotSession) 

91 def UpdateBotSession(self, request: UpdateBotSessionRequest, context: grpc.ServicerContext) -> BotSession: 

92 """Handles UpdateBotSessionRequest messages. 

93 

94 Args: 

95 request (UpdateBotSessionRequest): The incoming RPC request. 

96 context (grpc.ServicerContext): Context for the RPC call. 

97 """ 

98 LOGGER.debug(f"UpdateBotSession request from [{context.peer()}]") 

99 

100 if request.name != request.bot_session.name: 

101 raise InvalidArgumentError( 

102 "Name in UpdateBotSessionRequest does not match BotSession. " 

103 f" UpdateBotSessionRequest.name=[{request.name}] BotSession.name=[{request.bot_session.name}]" 

104 ) 

105 

106 instance_name = _instance_name_from_bot_name(request.name) 

107 instance = self.get_instance(instance_name) 

108 bot_session, metadata = instance.update_bot_session( 

109 request.bot_session, CancellationContext(context), deadline=context.time_remaining() 

110 ) 

111 

112 context.set_trailing_metadata(metadata) # type: ignore[arg-type] # tricky covariance issue. 

113 

114 return bot_session 

115 

116 @instanced(lambda r: _instance_name_from_bot_name(r.name)) 

117 @authorize 

118 @track_request_id 

119 def PostBotEventTemp(self, request: PostBotEventTempRequest, context: grpc.ServicerContext) -> empty_pb2.Empty: 

120 """Handles PostBotEventTempRequest messages. 

121 

122 Args: 

123 request (PostBotEventTempRequest): The incoming RPC request. 

124 context (grpc.ServicerContext): Context for the RPC call. 

125 """ 

126 LOGGER.info(f"PostBotEventTemp request from [{context.peer()}]") 

127 

128 context.set_code(grpc.StatusCode.UNIMPLEMENTED) 

129 

130 return empty_pb2.Empty() 

131 

132 def query_n_bots(self) -> int: 

133 return sum(map(self.query_n_bots_for_instance, self.instances)) 

134 

135 def query_n_bots_for_instance(self, instance_name: str) -> int: 

136 if instance := self.instances.get(instance_name): 

137 return instance.count_bots() 

138 return 0 

139 

140 def query_n_bots_for_status(self, bot_status: BotStatus) -> int: 

141 return sum(self.query_n_bots_for_instance_for_status(instance, bot_status) for instance in self.instances) 

142 

143 def query_n_bots_for_instance_for_status(self, instance_name: str, bot_status: BotStatus) -> int: 

144 if instance := self.instances.get(instance_name): 

145 return instance.count_bots_by_status(bot_status) 

146 return 0