1 # Copyright 2017 The Chromium OS Authors. All rights reserved. 2 # Use of this source code is governed by a BSD-style license that can be 3 # found in the LICENSE file. 4 5 import logging 6 7 import common 8 from autotest_lib.frontend.afe.json_rpc import proxy as rpc_proxy 9 from autotest_lib.server.hosts import host_info 10 from autotest_lib.server.cros.dynamic_suite import frontend_wrappers 11 12 class AfeStore(host_info.CachingHostInfoStore): 13 """Directly interact with the (given) AFE for host information.""" 14 15 _RETRYING_AFE_TIMEOUT_MIN = 5 16 _RETRYING_AFE_RETRY_DELAY_SEC = 10 17 18 def __init__(self, hostname, afe=None): 19 """ 20 @param hostname: The name of the host for which we want to track host 21 information. 22 @param afe: A frontend.AFE object to make RPC calls. Will create one 23 internally if None. 24 """ 25 super(AfeStore, self).__init__() 26 self._hostname = hostname 27 self._afe = afe 28 if self._afe is None: 29 self._afe = frontend_wrappers.RetryingAFE( 30 timeout_min=self._RETRYING_AFE_TIMEOUT_MIN, 31 delay_sec=self._RETRYING_AFE_RETRY_DELAY_SEC) 32 33 34 def __str__(self): 35 return '%s[%s]' % (type(self).__name__, self._hostname) 36 37 38 def _refresh_impl(self): 39 """Obtains HostInfo directly from the AFE.""" 40 try: 41 hosts = self._afe.get_hosts(hostname=self._hostname) 42 except rpc_proxy.JSONRPCException as e: 43 raise host_info.StoreError(e) 44 45 if not hosts: 46 raise host_info.StoreError('No hosts founds with hostname: %s' % 47 self._hostname) 48 49 if len(hosts) > 1: 50 logging.warning( 51 'Found %d hosts with the name %s. Picking the first one.', 52 len(hosts), self._hostname) 53 host = hosts[0] 54 return host_info.HostInfo(host.labels, host.attributes) 55 56 57 def _commit_impl(self, new_info): 58 """Commits HostInfo back to the AFE. 59 60 @param new_info: The new HostInfo to commit. 61 """ 62 # TODO(pprabhu) crbug.com/680322 63 # This method has a potentially malignent race condition. We obtain a 64 # copy of HostInfo from the AFE and then add/remove labels / attribtes 65 # based on that. If another user tries to commit it's changes in 66 # parallel, we'll end up with corrupted labels / attributes. 67 old_info = self._refresh_impl() 68 self._remove_labels_on_afe( 69 list(set(old_info.labels) - set(new_info.labels))) 70 self._add_labels_on_afe( 71 list(set(new_info.labels) - set(old_info.labels))) 72 self._update_attributes_on_afe(old_info.attributes, new_info.attributes) 73 74 75 def _remove_labels_on_afe(self, labels): 76 """Requests the AFE to remove the given labels. 77 78 @param labels: Remove these. 79 """ 80 if not labels: 81 return 82 83 logging.debug('removing labels: %s', labels) 84 try: 85 self._afe.run('host_remove_labels', id=self._hostname, 86 labels=labels) 87 except rpc_proxy.JSONRPCException as e: 88 raise host_info.StoreError(e) 89 90 91 def _add_labels_on_afe(self, labels): 92 """Requests the AFE to add the given labels. 93 94 @param labels: Add these. 95 """ 96 if not labels: 97 return 98 99 logging.info('adding labels: %s', labels) 100 try: 101 self._afe.run('host_add_labels', id=self._hostname, labels=labels) 102 except rpc_proxy.JSONRPCException as e: 103 raise host_info.StoreError(e) 104 105 106 def _update_attributes_on_afe(self, old_attributes, new_attributes): 107 """Updates host attributes on the afe to give dict. 108 109 @param old_attributes: The current attributes on AFE. 110 @param new_attributes: The new host attributes dict to set to. 111 """ 112 left_only, right_only, differing = _dict_diff(old_attributes, 113 new_attributes) 114 for key in left_only: 115 self._afe.set_host_attribute(key, None, hostname=self._hostname) 116 for key in right_only | differing: 117 self._afe.set_host_attribute(key, new_attributes[key], 118 hostname=self._hostname) 119 120 121 def _dict_diff(left_dict, right_dict): 122 """Return the keys where the given dictionaries differ. 123 124 This function assumes that the values in the dictionary support checking for 125 equality. 126 127 @param left_dict: The "left" dictionary in the diff. 128 @param right_dict: The "right" dictionary in the diff. 129 @returns: A 3-tuple (left_only, right_only, differing) of keys where 130 left_only contains the keys that exist in left_dict only, right_only 131 contains keys that exist in right_dict only and differing contains 132 keys that exist in both, but where values differ. 133 """ 134 left_keys = set(left_dict) 135 right_keys = set(right_dict) 136 differing_keys = {key for key in left_keys & right_keys 137 if left_dict[key] != right_dict[key]} 138 return left_keys - right_keys, right_keys - left_keys, differing_keys 139
