Coverage for /builds/BuildGrid/buildgrid/buildgrid/bot/session.py: 70.59%

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

17Bot Session 

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

19 

20Allows connections 

21""" 

22import asyncio 

23import logging 

24import os 

25import platform 

26 

27import grpc 

28 

29from buildgrid._enums import BotStatus, LeaseState 

30from buildgrid._exceptions import FailedPreconditionError 

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

32from buildgrid._protos.google.rpc import code_pb2 

33from .tenantmanager import TenantManager 

34 

35 

36class BotSession: 

37 def __init__(self, instance_name, bots_interface, hardware_interface, work, 

38 context=None): 

39 """ Unique bot ID within the farm used to identify this bot 

40 Needs to be human readable. 

41 All prior sessions with bot_id of same ID are invalidated. 

42 If a bot attempts to update an invalid session, it must be rejected and 

43 may be put in quarantine. 

44 """ 

45 self.__logger = logging.getLogger(__name__) 

46 

47 self._bots_interface = bots_interface 

48 self._hardware_interface = hardware_interface 

49 

50 self._status = BotStatus.OK.value 

51 self._tenant_manager = TenantManager() 

52 

53 self.__instance_name = instance_name 

54 self.__bot_id = f'{instance_name}.{platform.node()}.{os.getpid()}' 

55 self.__name = None 

56 

57 self._work = work 

58 self._context = context 

59 

60 self.__connected = False 

61 

62 @property 

63 def bot_id(self): 

64 return self.__bot_id 

65 

66 @property 

67 def connected(self): 

68 return self.__connected 

69 

70 async def run(self): 

71 interval = self._bots_interface.interval 

72 executing_interval = self._bots_interface.executing_interval 

73 try: 

74 while True: 

75 if not self.connected: 

76 self.create_bot_session() 

77 else: 

78 self.update_bot_session() 

79 

80 if not self.connected: 

81 await asyncio.sleep(interval) 

82 else: 

83 await self._tenant_manager.wait_on_tenants(executing_interval) 

84 except asyncio.CancelledError: 

85 pass 

86 

87 def create_bot_session(self): 

88 self.__logger.debug("Creating bot session") 

89 

90 session = self._bots_interface.create_bot_session(self.__instance_name, self.get_pb2()) 

91 if session in list(grpc.StatusCode): 

92 self.__connected = False 

93 return 

94 self.__connected = True 

95 self.__name = session.name 

96 

97 self.__logger.info(f"Created bot session with name: [{self.__name}]") 

98 

99 for lease in session.leases: 

100 self._register_lease(lease) 

101 

102 def update_bot_session(self): 

103 self.__logger.debug(f"Updating bot session: [{self.__bot_id}]") 

104 

105 session = self._bots_interface.update_bot_session(self.get_pb2()) 

106 if session in list(grpc.StatusCode): 

107 self.__connected = False 

108 return 

109 self.__connected = True 

110 server_ids = [] 

111 

112 for lease in session.leases: 

113 server_ids.append(lease.id) 

114 

115 lease_state = LeaseState(lease.state) 

116 if lease_state == LeaseState.PENDING: 

117 self._register_lease(lease) 

118 

119 elif lease_state == LeaseState.CANCELLED: 

120 self._tenant_manager.cancel_tenancy(lease.id) 

121 

122 closed_lease_ids = [x for x in self._tenant_manager.get_lease_ids() if x not in server_ids] 

123 

124 for lease_id in closed_lease_ids: 

125 self._tenant_manager.cancel_tenancy(lease_id) 

126 self._tenant_manager.remove_tenant(lease_id) 

127 

128 def get_pb2(self): 

129 return bots_pb2.BotSession(worker=self._hardware_interface.get_worker_pb2(), 

130 status=self._status, 

131 leases=self._tenant_manager.get_leases(), 

132 bot_id=self.__bot_id, 

133 name=self.__name) 

134 

135 def _register_lease(self, lease): 

136 lease_id = lease.id 

137 try: 

138 self._tenant_manager.create_tenancy(lease) 

139 

140 except KeyError as e: 

141 self.__logger.error(e) 

142 

143 else: 

144 try: 

145 self._hardware_interface.configure_hardware(lease.requirements) 

146 

147 except FailedPreconditionError as e: 

148 self.__logger.error(e) 

149 self._tenant_manager.complete_lease(lease_id, status=code_pb2.FailedPreconditionError) 

150 

151 else: 

152 self._tenant_manager.create_work(lease_id, self._work, self._context)