• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/python
2
3# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6"""A small wrapper script, iterates through
7the known hosts and tries to call get_labels()
8to discover host functionality, and adds these
9detected labels to host.
10
11Limitations:
12 - Does not keep a count of how many labels were
13   actually added.
14 - If a label is added by this script because it
15   is detected as supported by get_labels, but later becomes
16   unsupported, this script has no way to know that it
17   should be removed, so it will remain attached to the host.
18   See crosbug.com/38569
19"""
20
21
22from multiprocessing import pool
23import logging
24import socket
25import argparse
26import sys
27
28import common
29
30from autotest_lib.server import hosts
31from autotest_lib.server import frontend
32from autotest_lib.client.common_lib import error
33
34
35# A list of label prefix that each dut should only have one of such label with
36# the given prefix, e.g., a dut can't have both labels of power:battery and
37# power:AC_only.
38SINGLETON_LABEL_PREFIX = ['power:']
39
40def add_missing_labels(afe, hostname):
41    """
42    Queries the detectable labels supported by the given host,
43    and adds those labels to the host.
44
45    @param afe: A frontend.AFE() instance.
46    @param hostname: The host to query and update.
47
48    @return: True on success.
49             False on failure to fetch labels or to add any individual label.
50    """
51    host = None
52    try:
53        host = hosts.create_host(hostname)
54        labels = host.get_labels()
55    except socket.gaierror:
56        logging.warning('Unable to establish ssh connection to hostname '
57                        '%s. Skipping.', hostname)
58        return False
59    except error.AutoservError:
60        logging.warning('Unable to query labels on hostname %s. Skipping.',
61                         hostname)
62        return False
63    finally:
64        if host:
65            host.close()
66
67    afe_host = afe.get_hosts(hostname=hostname)[0]
68
69    label_matches = afe.get_labels(name__in=labels)
70
71    for label in label_matches:
72        singleton_prefixes = [p for p in SINGLETON_LABEL_PREFIX
73                              if label.name.startswith(p)]
74        if len(singleton_prefixes) == 1:
75            singleton_prefix = singleton_prefixes[0]
76            # Delete existing label with `singleton_prefix`
77            labels_to_delete = [l for l in afe_host.labels
78                                if l.startswith(singleton_prefix) and
79                                not l in labels]
80            if labels_to_delete:
81                logging.warning('Removing label %s', labels_to_delete)
82                afe_labels_to_delete = afe.get_labels(name__in=labels_to_delete)
83                for afe_label in afe_labels_to_delete:
84                    afe_label.remove_hosts(hosts=[hostname])
85        label.add_hosts(hosts=[hostname])
86
87    missing_labels = set(labels) - set([l.name for l in label_matches])
88
89    if missing_labels:
90        for label in missing_labels:
91            logging.warning('Unable to add label %s to host %s. '
92                            'Skipping unknown label.', label, hostname)
93        return False
94
95    return True
96
97
98def main():
99    """"
100    Entry point for add_detected_host_labels script.
101    """
102
103    parser = argparse.ArgumentParser()
104    parser.add_argument('-s', '--silent', dest='silent', action='store_true',
105                        help='Suppress all but critical logging messages.')
106    parser.add_argument('-i', '--info', dest='info_only', action='store_true',
107                        help='Suppress logging messages below INFO priority.')
108    parser.add_argument('-m', '--machines', dest='machines',
109                        help='Comma separated list of machines to check.')
110    options = parser.parse_args()
111
112    if options.silent and options.info_only:
113        print 'The -i and -s flags cannot be used together.'
114        parser.print_help()
115        return 0
116
117
118    if options.silent:
119        logging.disable(logging.CRITICAL)
120
121    if options.info_only:
122        logging.disable(logging.DEBUG)
123
124    threadpool = pool.ThreadPool()
125    afe = frontend.AFE()
126
127    if options.machines:
128        hostnames = [m.strip() for m in options.machines.split(',')]
129    else:
130        hostnames = afe.get_hostnames()
131    successes = sum(threadpool.imap_unordered(
132                        lambda x: add_missing_labels(afe, x),
133                        hostnames))
134    attempts = len(hostnames)
135
136    logging.info('Label updating finished. Failed update on %d out of %d '
137                 'hosts.', attempts-successes, attempts)
138
139    return 0
140
141
142if __name__ == '__main__':
143    sys.exit(main())
144