1# Copyright (c) 2012 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
6#pylint: disable-msg=C0111
7def order_by_complexity(host_spec_list):
8    """
9    Returns a new list of HostSpecs, ordered from most to least complex.
10
11    Currently, 'complex' means that the spec contains more labels.
12    We may want to get smarter about this.
13
14    @param host_spec_list: a list of HostSpec objects.
15    @return a new list of HostSpec, ordered from most to least complex.
16    """
17    def extract_label_list_len(host_spec):
18        return len(host_spec.labels)
19    return sorted(host_spec_list, key=extract_label_list_len, reverse=True)
20
21
22def is_simple_list(host_spec_list):
23    """
24    Returns true if this is a 'simple' list of HostSpec objects.
25
26    A 'simple' list of HostSpec objects is defined as a list of one HostSpec.
27
28    @param host_spec_list: a list of HostSpec objects.
29    @return True if this is a list of size 1, False otherwise.
30    """
31    return len(host_spec_list) == 1
32
33
34def simple_get_spec_and_hosts(host_specs, hosts_per_spec):
35    """Given a simple list of HostSpec, extract hosts from hosts_per_spec.
36
37    Given a simple list of HostSpec objects, pull out the spec and use it to
38    get the associated hosts out of hosts_per_spec.  Return the spec and the
39    host list as a pair.
40
41    @param host_specs: an iterable of HostSpec objects.
42    @param hosts_per_spec: map of {HostSpec: [list, of, hosts]}
43    @return (HostSpec, [list, of, hosts]}
44    """
45    spec = host_specs.pop()
46    return spec, hosts_per_spec[spec]
47
48
49class HostGroup(object):
50    """A high-level specification of a group of hosts.
51
52    A HostGroup represents a group of hosts against which a job can be
53    scheduled.  An instance is capable of returning arguments that can specify
54    this group in a call to AFE.create_job().
55    """
56    def __init__(self):
57        pass
58
59
60    def as_args(self):
61        """Return args suitable for passing to AFE.create_job()."""
62        raise NotImplementedError()
63
64
65    def size(self):
66        """Returns the number of hosts specified by the group."""
67        raise NotImplementedError()
68
69
70    def mark_host_success(self, hostname):
71        """Marks the provided host as successfully reimaged.
72
73        @param hostname: the name of the host that was reimaged.
74        """
75        raise NotImplementedError()
76
77
78    def enough_hosts_succeeded(self):
79        """Returns True if enough hosts in the group were reimaged for use."""
80        raise NotImplementedError()
81
82
83    #pylint: disable-msg=C0111
84    @property
85    def unsatisfied_specs(self):
86        return []
87
88
89    #pylint: disable-msg=C0111
90    @property
91    def doomed_specs(self):
92        return []
93
94
95class ExplicitHostGroup(HostGroup):
96    """A group of hosts, specified by name, to be reimaged for use.
97
98    @var _hostname_data_dict: {hostname: HostData()}.
99    """
100
101    class HostData(object):
102        """A HostSpec of a given host, and whether it reimaged successfully."""
103        def __init__(self, spec):
104            self.spec = spec
105            self.image_success = False
106
107
108    def __init__(self, hosts_per_spec={}):
109        """Constructor.
110
111        @param hosts_per_spec: {HostSpec: [list, of, hosts]}.
112                               Each host can appear only once.
113        """
114        self._hostname_data_dict = {}
115        self._potentially_unsatisfied_specs = []
116        for spec, host_list in hosts_per_spec.iteritems():
117            for host in host_list:
118                self.add_host_for_spec(spec, host)
119
120
121    def _get_host_datas(self):
122        return self._hostname_data_dict.itervalues()
123
124
125    def as_args(self):
126        return {'hosts': self._hostname_data_dict.keys()}
127
128
129    def size(self):
130        return len(self._hostname_data_dict)
131
132
133    def mark_host_success(self, hostname):
134        self._hostname_data_dict[hostname].image_success = True
135
136
137    def enough_hosts_succeeded(self):
138        """If _any_ hosts were reimaged, that's enough."""
139        return True in [d.image_success for d in self._get_host_datas()]
140
141
142    def add_host_for_spec(self, spec, host):
143        """Add a new host for the given HostSpec to the group.
144
145        @param spec: HostSpec to associate host with.
146        @param host: a Host object; each host can appear only once.
147                     If None, this spec will be relegated to the list of
148                     potentially unsatisfied specs.
149        """
150        if not host:
151            if spec not in [d.spec for d in self._get_host_datas()]:
152                self._potentially_unsatisfied_specs.append(spec)
153            return
154
155        if self.contains_host(host):
156            raise ValueError('A Host can appear in an '
157                             'ExplicitHostGroup only once.')
158        if spec in self._potentially_unsatisfied_specs:
159            self._potentially_unsatisfied_specs.remove(spec)
160        self._hostname_data_dict[host.hostname] = self.HostData(spec)
161
162
163    def contains_host(self, host):
164        """Whether host is already part of this HostGroup
165
166        @param host: a Host object.
167        @return True if the host is already tracked; False otherwise.
168        """
169        return host.hostname in self._hostname_data_dict
170
171
172    @property
173    def unsatisfied_specs(self):
174        unsatisfied = []
175        for spec in self._potentially_unsatisfied_specs:
176            # If a spec in _potentially_unsatisfied_specs is a subset of some
177            # satisfied spec, then it's not unsatisfied.
178            if filter(lambda d: spec.is_subset(d.spec), self._get_host_datas()):
179                continue
180            unsatisfied.append(spec)
181        return unsatisfied
182
183
184    @property
185    def doomed_specs(self):
186        ok = set()
187        possibly_doomed = set()
188        for data in self._get_host_datas():
189            # If imaging succeeded for any host that satisfies a spec,
190            # it's definitely not doomed.
191            if data.image_success:
192                ok.add(data.spec)
193            else:
194                possibly_doomed.add(data.spec)
195        # If a spec is not a subset of any ok spec, it's doomed.
196        return set([s for s in possibly_doomed if not filter(s.is_subset, ok)])
197
198
199class MetaHostGroup(HostGroup):
200    """A group of hosts, specified by a meta_host and deps, to be reimaged.
201
202    @var _meta_hosts: a meta_host, as expected by AFE.create_job()
203    @var _dependencies: list of dependencies that all hosts to be used
204                        must satisfy
205    @var _successful_hosts: set of successful hosts.
206    """
207    def __init__(self, labels, num):
208        """Constructor.
209
210        Given a set of labels specifying what kind of hosts we need,
211        and the num of hosts we need, build a meta_host and dependency list
212        that represent this group of hosts.
213
214        @param labels: list of labels indicating what kind of hosts need
215                       to be reimaged.
216        @param num: how many hosts we'd like to reimage.
217        """
218        self._spec = HostSpec(labels)
219        self._meta_hosts = labels[:1]*num
220        self._dependencies = labels[1:]
221        self._successful_hosts = set()
222
223
224    def as_args(self):
225        return {'meta_hosts': self._meta_hosts,
226                'dependencies': self._dependencies}
227
228
229    def size(self):
230        return len(self._meta_hosts)
231
232
233    def mark_host_success(self, hostname):
234        self._successful_hosts.add(hostname)
235
236
237    def enough_hosts_succeeded(self):
238        return self._successful_hosts
239
240
241    @property
242    def doomed_specs(self):
243        if self._successful_hosts:
244            return []
245        return [self._spec]
246
247
248def _safeunion(iter_a, iter_b):
249    """Returns an immutable set that contains the union of two iterables.
250
251    This function returns a frozen set containing the all the elements of
252    two iterables, regardless of whether those iterables are lists, sets,
253    or whatever.
254
255    @param iter_a: The first iterable.
256    @param iter_b: The second iterable.
257    @returns: An immutable union of the contents of iter_a and iter_b.
258    """
259    return frozenset({a for a in iter_a} | {b for b in iter_b})
260
261
262
263class HostSpec(object):
264    """Specifies a kind of host on which dependency-having tests can be run.
265
266    Wraps a list of labels, for the purposes of specifying a set of hosts
267    on which a test with matching dependencies can be run.
268    """
269
270    def __init__(self, base, extended=[]):
271        self._labels = _safeunion(base, extended)
272        # To amortize cost of __hash__()
273        self._str = 'HostSpec %r' % sorted(self._labels)
274        self._trivial = extended == []
275
276
277    #pylint: disable-msg=C0111
278    @property
279    def labels(self):
280        # Can I just do this as a set?  Inquiring minds want to know.
281        return sorted(self._labels)
282
283
284    #pylint: disable-msg=C0111
285    @property
286    def is_trivial(self):
287        return self._trivial
288
289
290    #pylint: disable-msg=C0111
291    def is_subset(self, other):
292        return self._labels <= other._labels
293
294
295    def __str__(self):
296        return self._str
297
298
299    def __repr__(self):
300        return self._str
301
302
303    def __lt__(self, other):
304        return str(self) < str(other)
305
306
307    def __le__(self, other):
308        return str(self) <= str(other)
309
310
311    def __eq__(self, other):
312        return str(self) == str(other)
313
314
315    def __ne__(self, other):
316        return str(self) != str(other)
317
318
319    def __gt__(self, other):
320        return str(self) > str(other)
321
322
323    def __ge__(self, other):
324        return str(self) >= str(other)
325
326
327    def __hash__(self):
328        """Allows instances to be correctly deduped when used in a set."""
329        return hash(str(self))
330