1#
2# Copyright 2008 Google Inc. All Rights Reserved.
3
4"""
5The shard module contains the objects and methods used to
6manage shards in Autotest.
7
8The valid actions are:
9create:      creates shard
10remove:      deletes shard(s)
11list:        lists shards with label
12
13See topic_common.py for a High Level Design and Algorithm.
14"""
15
16import sys
17from autotest_lib.cli import topic_common, action_common
18
19
20class shard(topic_common.atest):
21    """shard class
22    atest shard [create|delete|list] <options>"""
23    usage_action = '[create|delete|list]'
24    topic = msg_topic = 'shard'
25    msg_items = '<shards>'
26
27    def __init__(self):
28        """Add to the parser the options common to all the
29        shard actions"""
30        super(shard, self).__init__()
31
32        self.topic_parse_info = topic_common.item_parse_info(
33            attribute_name='shards',
34            use_leftover=True)
35
36
37    def get_items(self):
38        return self.shards
39
40
41class shard_help(shard):
42    """Just here to get the atest logic working.
43    Usage is set by its parent"""
44    pass
45
46
47class shard_list(action_common.atest_list, shard):
48    """Class for running atest shard list"""
49
50    def execute(self):
51        filters = {}
52        if self.shards:
53            filters['hostname__in'] = self.shards
54        return super(shard_list, self).execute(op='get_shards',
55                                               filters=filters)
56
57
58    def warn_if_label_assigned_to_multiple_shards(self, results):
59        """Prints a warning if one label is assigned to multiple shards.
60
61        This should never happen, but if it does, better be safe.
62
63        @param results: Results as passed to output().
64        """
65        assigned_labels = set()
66        for line in results:
67            for label in line['labels']:
68                if label in assigned_labels:
69                    sys.stderr.write('WARNING: label %s is assigned to '
70                                     'multiple shards.\n'
71                                     'This will lead to unpredictable behavor '
72                                     'in which hosts and jobs will be assigned '
73                                     'to which shard.\n')
74                assigned_labels.add(label)
75
76
77    def output(self, results):
78        self.warn_if_label_assigned_to_multiple_shards(results)
79        super(shard_list, self).output(results, ['hostname', 'labels'])
80
81
82class shard_create(action_common.atest_create, shard):
83    """Class for running atest shard create -l <label> <shard>"""
84    def __init__(self):
85        super(shard_create, self).__init__()
86        self.parser.add_option('-l', '--labels',
87                               help=('Assign LABELs to the SHARD. All jobs that '
88                                     'require one of the labels will be run on  '
89                                     'the shard. List multiple labels separated '
90                                     'by a comma.'),
91                               type='string',
92                               metavar='LABELS')
93
94
95    def parse(self):
96        (options, leftover) = super(shard_create, self).parse(
97                req_items='shards')
98        if not options.labels:
99            print ('Must provide one or more labels separated by a comma '
100                  'with -l <labels>')
101            self.parser.print_help()
102            sys.exit(1)
103        self.data_item_key = 'hostname'
104        self.data['labels'] = options.labels
105        return (options, leftover)
106
107
108class shard_delete(action_common.atest_delete, shard):
109    """Class for running atest shard delete <shards>"""
110
111    def parse(self):
112        (options, leftover) = super(shard_delete, self).parse()
113        self.data_item_key = 'hostname'
114        return (options, leftover)
115
116
117    def execute(self, *args, **kwargs):
118        print 'Please ensure the shard host is powered off.'
119        print ('Otherwise DUTs might be used by multiple shards at the same '
120               'time, which will lead to serious correctness problems.')
121        return super(shard_delete, self).execute(*args, **kwargs)
122