xref: /aosp_15_r20/external/autotest/server/system_utils.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7) 
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
16
17
18def get_drones():
19    """Get a list of drones from server database or global config.
20    """
21    return []
22
23
24def get_shards():
25    """Get a list of shards from server database or global config.
26    """
27    config = global_config.global_config
28    shards = config.get_config_value('SERVER', 'shards', default='')
29    return [hostname.strip() for hostname in shards.split(',')]
30
31
32class DroneCache(object):
33    """A cache object to store drone list related data.
34
35    The cache is added to avoid repeated calls from each agent task. The cache
36    should be refreshed every tick.
37    """
38
39    # A cache of unrestricted drones.
40    unrestricted_drones = None
41
42    # A cache of a dict of (drone, ip).
43    drone_ip_map = None
44
45    @classmethod
46    def refresh(cls, restricted_subnets=utils.RESTRICTED_SUBNETS):
47        """Refresh the cache.
48
49        @param restricted_subnets: A list of restricted subnet, default is set
50                to restricted_subnets in global config.
51        """
52        new_drone_ip_map = {}
53        new_unrestricted_drones = []
54        for drone in get_drones():
55            new_drone_ip_map[drone] = utils.get_ip_address(drone)
56            if (not restricted_subnets or
57                not utils.get_restricted_subnet(new_drone_ip_map[drone],
58                                                restricted_subnets)):
59                new_unrestricted_drones.append(drone)
60        cls.drone_ip_map = new_drone_ip_map
61        cls.unrestricted_drones = new_unrestricted_drones
62
63
64    @classmethod
65    def get_unrestricted_drones(
66                cls, restricted_subnets=utils.RESTRICTED_SUBNETS):
67        """Get a list of cached unrestricted drones.
68
69        @param restricted_subnets: A list of restricted subnet, default is set
70                to restricted_subnets in global config.
71        """
72        if not cls.unrestricted_drones:
73            cls.refresh(restricted_subnets)
74
75        return copy.copy(cls.unrestricted_drones)
76
77
78    @classmethod
79    def get_drone_ip_map(cls, restricted_subnets=utils.RESTRICTED_SUBNETS):
80        """Get a dict of (drone, ip).
81
82        @param restricted_subnets: A list of restricted subnet, default is set
83                to restricted_subnets in global config.
84        """
85        if not cls.drone_ip_map:
86            cls.refresh(restricted_subnets)
87
88        return copy.copy(cls.drone_ip_map)
89