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
12add_boards:  add boards to a given shard
13
14See topic_common.py for a High Level Design and Algorithm.
15"""
16
17import sys
18from autotest_lib.cli import topic_common, action_common
19
20
21class shard(topic_common.atest):
22    """shard class
23    atest shard [create|delete|list|add_boards] <options>"""
24    usage_action = '[create|delete|list|add_boards]'
25    topic = msg_topic = 'shard'
26    msg_items = '<shards>'
27
28    def __init__(self):
29        """Add to the parser the options common to all the
30        shard actions"""
31        super(shard, self).__init__()
32
33        self.topic_parse_info = topic_common.item_parse_info(
34            attribute_name='shards',
35            use_leftover=True)
36
37
38    def get_items(self):
39        return self.shards
40
41
42class shard_help(shard):
43    """Just here to get the atest logic working.
44    Usage is set by its parent"""
45    pass
46
47
48class shard_list(action_common.atest_list, shard):
49    """Class for running atest shard list"""
50
51    def execute(self):
52        filters = {}
53        if self.shards:
54            filters['hostname__in'] = self.shards
55        return super(shard_list, self).execute(op='get_shards',
56                                               filters=filters)
57
58
59    def warn_if_label_assigned_to_multiple_shards(self, results):
60        """Prints a warning if one label is assigned to multiple shards.
61
62        This should never happen, but if it does, better be safe.
63
64        @param results: Results as passed to output().
65        """
66        assigned_labels = set()
67        for line in results:
68            for label in line['labels']:
69                if label in assigned_labels:
70                    sys.stderr.write('WARNING: label %s is assigned to '
71                                     'multiple shards.\n'
72                                     'This will lead to unpredictable behavor '
73                                     'in which hosts and jobs will be assigned '
74                                     'to which shard.\n')
75                assigned_labels.add(label)
76
77
78    def output(self, results):
79        self.warn_if_label_assigned_to_multiple_shards(results)
80        super(shard_list, self).output(results, ['hostname', 'labels'])
81
82
83class shard_create(action_common.atest_create, shard):
84    """Class for running atest shard create -l <label> <shard>"""
85    def __init__(self):
86        super(shard_create, self).__init__()
87        self.parser.add_option('-l', '--labels',
88                               help=('Assign LABELs to the SHARD. All jobs that '
89                                     'require one of the labels will be run on  '
90                                     'the shard. List multiple labels separated '
91                                     'by a comma.'),
92                               type='string',
93                               metavar='LABELS')
94
95
96    def parse(self):
97        (options, leftover) = super(shard_create, self).parse(
98                req_items='shards')
99        self.data_item_key = 'hostname'
100        self.data['labels'] = options.labels or ''
101        return (options, leftover)
102
103
104class shard_add_boards(shard_create):
105    """Class for running atest shard add_boards -l <label> <shard>"""
106    usage_action = 'add_boards'
107    op_action = 'add_boards'
108    msg_done = 'Added boards for'
109
110    def execute(self):
111        """Running the rpc to add boards to the target shard.
112
113        Returns:
114          A tuple, 1st element is the target shard. 2nd element is the list of
115          boards labels to be added to the shard.
116        """
117        target_shard = self.shards[0]
118        self.data[self.data_item_key] = target_shard
119        super(shard_add_boards, self).execute_rpc(op='add_board_to_shard',
120                                                  item=target_shard,
121                                                  **self.data)
122        return (target_shard, self.data['labels'])
123
124
125class shard_delete(action_common.atest_delete, shard):
126    """Class for running atest shard delete <shards>"""
127
128    def parse(self):
129        (options, leftover) = super(shard_delete, self).parse()
130        self.data_item_key = 'hostname'
131        return (options, leftover)
132
133
134    def execute(self, *args, **kwargs):
135        print 'Please ensure the shard host is powered off.'
136        print ('Otherwise DUTs might be used by multiple shards at the same '
137               'time, which will lead to serious correctness problems.')
138        return super(shard_delete, self).execute(*args, **kwargs)
139