1# Copyright 2016 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"""This module provides utility functions to get the servers used by Autotest 6system, from server database or global config. The standalone module is needed 7to avoid circular imports. 8 9""" 10 11import copy 12 13import common 14from autotest_lib.client.common_lib import global_config 15from autotest_lib.client.common_lib import utils 16from autotest_lib.site_utils import server_manager_utils 17 18 19def get_drones(): 20 """Get a list of drones from server database or global config. 21 """ 22 if server_manager_utils.use_server_db(): 23 return server_manager_utils.get_drones() 24 else: 25 return [] 26 27 28def get_shards(): 29 """Get a list of shards from server database or global config. 30 """ 31 if server_manager_utils.use_server_db(): 32 return server_manager_utils.get_shards() 33 else: 34 config = global_config.global_config 35 shards = config.get_config_value( 36 'SERVER', 'shards', default='') 37 return [hostname.strip() for hostname in shards.split(',')] 38 39 40class DroneCache(object): 41 """A cache object to store drone list related data. 42 43 The cache is added to avoid repeated calls from each agent task. The cache 44 should be refreshed every tick. 45 """ 46 47 # A cache of unrestricted drones. 48 unrestricted_drones = None 49 50 # A cache of a dict of (drone, ip). 51 drone_ip_map = None 52 53 @classmethod 54 def refresh(cls, restricted_subnets=utils.RESTRICTED_SUBNETS): 55 """Refresh the cache. 56 57 @param restricted_subnets: A list of restricted subnet, default is set 58 to restricted_subnets in global config. 59 """ 60 new_drone_ip_map = {} 61 new_unrestricted_drones = [] 62 for drone in get_drones(): 63 new_drone_ip_map[drone] = utils.get_ip_address(drone) 64 if (not restricted_subnets or 65 not utils.get_restricted_subnet(new_drone_ip_map[drone], 66 restricted_subnets)): 67 new_unrestricted_drones.append(drone) 68 cls.drone_ip_map = new_drone_ip_map 69 cls.unrestricted_drones = new_unrestricted_drones 70 71 72 @classmethod 73 def get_unrestricted_drones( 74 cls, restricted_subnets=utils.RESTRICTED_SUBNETS): 75 """Get a list of cached unrestricted drones. 76 77 @param restricted_subnets: A list of restricted subnet, default is set 78 to restricted_subnets in global config. 79 """ 80 if not cls.unrestricted_drones: 81 cls.refresh(restricted_subnets) 82 83 return copy.copy(cls.unrestricted_drones) 84 85 86 @classmethod 87 def get_drone_ip_map(cls, restricted_subnets=utils.RESTRICTED_SUBNETS): 88 """Get a dict of (drone, ip). 89 90 @param restricted_subnets: A list of restricted subnet, default is set 91 to restricted_subnets in global config. 92 """ 93 if not cls.drone_ip_map: 94 cls.refresh(restricted_subnets) 95 96 return copy.copy(cls.drone_ip_map) 97