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