| OLD | NEW |
| 1 # Copyright 2016 The LUCI Authors. All rights reserved. | 1 # Copyright 2016 The LUCI Authors. All rights reserved. |
| 2 # Use of this source code is governed under the Apache License, Version 2.0 | 2 # Use of this source code is governed under the Apache License, Version 2.0 |
| 3 # that can be found in the LICENSE file. | 3 # that can be found in the LICENSE file. |
| 4 | 4 |
| 5 # This is a reimplementation of RemoteClientNative but it uses (will use) | 5 # This is a reimplementation of RemoteClientNative but it uses (will use) |
| 6 # a gRPC method to communicate with a server instead of REST. | 6 # a gRPC method to communicate with a server instead of REST. |
| 7 | 7 |
| 8 import json | 8 import json |
| 9 import logging | 9 import logging |
| 10 | 10 |
| 11 import grpc | 11 import grpc |
| 12 import google.protobuf.json_format | 12 import google.protobuf.json_format |
| 13 from proto_bot import swarming_bot_pb2 | 13 from proto_bot import swarming_bot_pb2 |
| 14 from remote_client_errors import InternalError | 14 from remote_client_errors import InternalError |
| 15 from remote_client_errors import PollError |
| 15 | 16 |
| 16 | 17 |
| 17 # How long to wait for a response from the server. Keeping the same as | 18 # How long to wait for a response from the server. Keeping the same as |
| 18 # the equivalent in remote_client.py for now. | 19 # the equivalent in remote_client.py for now. |
| 19 NET_CONNECTION_TIMEOUT_SEC = 5*60 | 20 NET_CONNECTION_TIMEOUT_SEC = 5*60 |
| 20 | 21 |
| 21 | 22 |
| 22 class RemoteClientGrpc(object): | 23 class RemoteClientGrpc(object): |
| 23 """RemoteClientGrpc knows how to make calls via gRPC. | 24 """RemoteClientGrpc knows how to make calls via gRPC. |
| 24 """ | 25 """ |
| (...skipping 80 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
| 105 d.name: d.values for d in response.bot_group_cfg.dimensions | 106 d.name: d.values for d in response.bot_group_cfg.dimensions |
| 106 }, | 107 }, |
| 107 }, | 108 }, |
| 108 } | 109 } |
| 109 logging.info('Completed handshake: %s', resp) | 110 logging.info('Completed handshake: %s', resp) |
| 110 return resp | 111 return resp |
| 111 | 112 |
| 112 def poll(self, attributes): | 113 def poll(self, attributes): |
| 113 request = swarming_bot_pb2.PollRequest() | 114 request = swarming_bot_pb2.PollRequest() |
| 114 self._attributes_json_to_proto(attributes, request.attributes) | 115 self._attributes_json_to_proto(attributes, request.attributes) |
| 115 # TODO(aludwin): gRPC-specific exception handling | 116 # TODO(aludwin): gRPC-specific exception handling (raise PollError). |
| 116 response = self._stub.Poll(request, timeout=NET_CONNECTION_TIMEOUT_SEC) | 117 response = self._stub.Poll(request, timeout=NET_CONNECTION_TIMEOUT_SEC) |
| 117 | 118 |
| 118 if response.cmd == swarming_bot_pb2.PollResponse.UPDATE: | 119 if response.cmd == swarming_bot_pb2.PollResponse.UPDATE: |
| 119 return 'update', response.version | 120 return 'update', response.version |
| 120 | 121 |
| 121 if response.cmd == swarming_bot_pb2.PollResponse.SLEEP: | 122 if response.cmd == swarming_bot_pb2.PollResponse.SLEEP: |
| 122 if not self._log_is_asleep: | 123 if not self._log_is_asleep: |
| 123 logging.info('Going to sleep') | 124 logging.info('Going to sleep') |
| 124 self._log_is_asleep = True | 125 self._log_is_asleep = True |
| 125 return 'sleep', response.sleep_time | 126 return 'sleep', response.sleep_time |
| (...skipping 33 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
| 159 # These keys are only needed by raw commands. While this method | 160 # These keys are only needed by raw commands. While this method |
| 160 # only supports isolated commands, the keys need to exist to avoid | 161 # only supports isolated commands, the keys need to exist to avoid |
| 161 # missing key errors. | 162 # missing key errors. |
| 162 'command': None, | 163 'command': None, |
| 163 'extra_args': None, | 164 'extra_args': None, |
| 164 } | 165 } |
| 165 logging.info('Received job manifest: %s', manifest) | 166 logging.info('Received job manifest: %s', manifest) |
| 166 self._log_is_asleep = False | 167 self._log_is_asleep = False |
| 167 return 'run', manifest | 168 return 'run', manifest |
| 168 | 169 |
| 169 raise ValueError('Unknown command in response: %s' % response) | 170 raise PollError('Unknown command in response: %s' % response) |
| 170 | 171 |
| 171 def get_bot_code(self, new_zip_fn, bot_version, _bot_id): | 172 def get_bot_code(self, new_zip_fn, bot_version, _bot_id): |
| 172 # TODO(aludwin): exception handling, pass bot_id | 173 # TODO(aludwin): exception handling, pass bot_id |
| 173 logging.info('Updating to version: %s', bot_version) | 174 logging.info('Updating to version: %s', bot_version) |
| 174 request = swarming_bot_pb2.BotUpdateRequest() | 175 request = swarming_bot_pb2.BotUpdateRequest() |
| 175 request.bot_version = bot_version | 176 request.bot_version = bot_version |
| 176 response = self._stub.BotUpdate(request, timeout=NET_CONNECTION_TIMEOUT_SEC) | 177 response = self._stub.BotUpdate(request, timeout=NET_CONNECTION_TIMEOUT_SEC) |
| 177 with open(new_zip_fn, 'wb') as f: | 178 with open(new_zip_fn, 'wb') as f: |
| 178 f.write(response.bot_code) | 179 f.write(response.bot_code) |
| 179 | 180 |
| (...skipping 34 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
| 214 def insert_dict_as_submessage(message, keyname, value): | 215 def insert_dict_as_submessage(message, keyname, value): |
| 215 """Encodes a dict as a Protobuf message. | 216 """Encodes a dict as a Protobuf message. |
| 216 | 217 |
| 217 The keyname for the message field is passed in to simplify the creation | 218 The keyname for the message field is passed in to simplify the creation |
| 218 of the submessage in the first place - you need to say getattr and not | 219 of the submessage in the first place - you need to say getattr and not |
| 219 simply refer to message.keyname since the former actually creates the | 220 simply refer to message.keyname since the former actually creates the |
| 220 submessage while the latter does not. | 221 submessage while the latter does not. |
| 221 """ | 222 """ |
| 222 sub_msg = getattr(message, keyname) | 223 sub_msg = getattr(message, keyname) |
| 223 google.protobuf.json_format.Parse(json.dumps(value), sub_msg) | 224 google.protobuf.json_format.Parse(json.dumps(value), sub_msg) |
| OLD | NEW |