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