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