lin
2025-07-30 fcd736bf35fd93b563e9bbf594f2aa7b62028cc9
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
 
"""RDB Host objects.
 
RDBHost: Basic host object, capable of retrieving fields of a host that
correspond to columns of the host table.
 
RDBServerHostWrapper: Server side host adapters that help in making a raw
database host object more ameanable to the classes and functions in the rdb
and/or rdb clients.
 
RDBClientHostWrapper: Scheduler host proxy that converts host information
returned by the rdb into a client host object capable of proxying updates
back to the rdb.
"""
 
import logging
 
from django.core import exceptions as django_exceptions
 
import common
from autotest_lib.client.common_lib import utils
from autotest_lib.frontend.afe import rdb_model_extensions as rdb_models
from autotest_lib.frontend.afe import models as afe_models
from autotest_lib.scheduler import rdb_requests
from autotest_lib.scheduler import rdb_utils
from autotest_lib.server import constants
from autotest_lib.utils import labellib
 
try:
    from chromite.lib import metrics
except ImportError:
    metrics = utils.metrics_mock
 
 
class RDBHost(object):
    """A python host object representing a django model for the host."""
 
    required_fields = set(
            rdb_models.AbstractHostModel.get_basic_field_names() + ['id'])
 
 
    def _update_attributes(self, new_attributes):
        """Updates attributes based on an input dictionary.
 
        Since reads are not proxied to the rdb this method caches updates to
        the host tables as class attributes.
 
        @param new_attributes: A dictionary of attributes to update.
        """
        for name, value in new_attributes.iteritems():
            setattr(self, name, value)
 
 
    def __init__(self, **kwargs):
        if self.required_fields - set(kwargs.keys()):
            raise rdb_utils.RDBException('Creating %s requires %s, got %s '
                    % (self.__class__, self.required_fields, kwargs.keys()))
        self._update_attributes(kwargs)
 
 
    @classmethod
    def get_required_fields_from_host(cls, host):
        """Returns all required attributes of the host parsed into a dict.
 
        Required attributes are defined as the attributes required to
        create an RDBHost, and mirror the columns of the host table.
 
        @param host: A host object containing all required fields as attributes.
        """
        required_fields_map = {}
        try:
            for field in cls.required_fields:
                required_fields_map[field] = getattr(host, field)
        except AttributeError as e:
            raise rdb_utils.RDBException('Required %s' % e)
        required_fields_map['id'] = host.id
        return required_fields_map
 
 
    def wire_format(self):
        """Returns information about this host object.
 
        @return: A dictionary of fields representing the host.
        """
        return RDBHost.get_required_fields_from_host(self)
 
 
class RDBServerHostWrapper(RDBHost):
    """A host wrapper for the base host object.
 
    This object contains all the attributes of the raw database columns,
    and a few more that make the task of host assignment easier. It handles
    the following duties:
        1. Serialization of the host object and foreign keys
        2. Conversion of label ids to label names, and retrieval of platform
        3. Checking the leased bit/status of a host before leasing it out.
    """
 
    def __init__(self, host):
        """Create an RDBServerHostWrapper.
 
        @param host: An instance of the Host model class.
        """
        host_fields = RDBHost.get_required_fields_from_host(host)
        super(RDBServerHostWrapper, self).__init__(**host_fields)
        self.labels = rdb_utils.LabelIterator(host.labels.all())
        self.acls = [aclgroup.id for aclgroup in host.aclgroup_set.all()]
        self.protection = host.protection
        platform = host.platform()
        # Platform needs to be a method, not an attribute, for
        # backwards compatibility with the rest of the host model.
        self.platform_name = platform.name if platform else None
        self.shard_id = host.shard_id
 
 
    def refresh(self, fields=None):
        """Refresh the attributes on this instance.
 
        @param fields: A list of fieldnames to refresh. If None
            all the required fields of the host are refreshed.
 
        @raises RDBException: If refreshing a field fails.
        """
        # TODO: This is mainly required for cache correctness. If it turns
        # into a bottleneck, cache host_ids instead of rdbhosts and rebuild
        # the hosts once before leasing them out. The important part is to not
        # trust the leased bit on a cached host.
        fields = self.required_fields if not fields else fields
        try:
            refreshed_fields = afe_models.Host.objects.filter(
                    id=self.id).values(*fields)[0]
        except django_exceptions.FieldError as e:
            raise rdb_utils.RDBException('Couldn\'t refresh fields %s: %s' %
                    fields, e)
        self._update_attributes(refreshed_fields)
 
 
    def lease(self):
        """Set the leased bit on the host object, and in the database.
 
        @raises RDBException: If the host is already leased.
        """
        self.refresh(fields=['leased'])
        if self.leased:
            raise rdb_utils.RDBException('Host %s is already leased' %
                                         self.hostname)
        self.leased = True
        # TODO: Avoid leaking django out of rdb.QueryManagers. This is still
        # preferable to calling save() on the host object because we're only
        # updating/refreshing a single indexed attribute, the leased bit.
        afe_models.Host.objects.filter(id=self.id).update(leased=self.leased)
 
 
    def wire_format(self, unwrap_foreign_keys=True):
        """Returns all information needed to scheduler jobs on the host.
 
        @param unwrap_foreign_keys: If true this method will retrieve and
            serialize foreign keys of the original host, which are stored
            in the RDBServerHostWrapper as iterators.
 
        @return: A dictionary of host information.
        """
        host_info = super(RDBServerHostWrapper, self).wire_format()
 
        if unwrap_foreign_keys:
            host_info['labels'] = self.labels.get_label_names()
            host_info['acls'] = self.acls
            host_info['platform_name'] = self.platform_name
            host_info['protection'] = self.protection
        return host_info
 
 
class RDBClientHostWrapper(RDBHost):
    """A client host wrapper for the base host object.
 
    This wrapper is used whenever the queue entry needs direct access
    to the host.
    """
    # Shows more detailed status of what a DUT is doing.
    _HOST_WORKING_METRIC = 'chromeos/autotest/dut_working'
    # Shows which hosts are working.
    _HOST_STATUS_METRIC = 'chromeos/autotest/dut_status'
    # Maps duts to pools.
    _HOST_POOL_METRIC = 'chromeos/autotest/dut_pool'
    # Shows which scheduler machines are using a DUT.
    _BOARD_SHARD_METRIC = 'chromeos/autotest/shard/board_presence'
 
 
    def __init__(self, **kwargs):
 
        # This class is designed to only check for the bare minimum
        # attributes on a host, so if a client tries accessing an
        # unpopulated foreign key it will result in an exception. Doing
        # so makes it easier to add fields to the rdb host without
        # updating all the clients.
        super(RDBClientHostWrapper, self).__init__(**kwargs)
 
        # TODO(beeps): Remove this once we transition to urls
        from autotest_lib.scheduler import rdb
        self.update_request_manager = rdb_requests.RDBRequestManager(
                rdb_requests.UpdateHostRequest, rdb.update_hosts)
        self.dbg_str = ''
        self.metadata = {}
        # We access labels for metrics generation below and it's awkward not
        # knowing if labels were populated or not.
        if not hasattr(self, 'labels'):
            self.labels = ()
 
 
 
    def _update(self, payload):
        """Send an update to rdb, save the attributes of the payload locally.
 
        @param: A dictionary representing 'key':value of the update required.
 
        @raises RDBException: If the update fails.
        """
        logging.info('Host %s in %s updating %s through rdb on behalf of: %s ',
                     self.hostname, self.status, payload, self.dbg_str)
        self.update_request_manager.add_request(host_id=self.id,
                payload=payload)
        for response in self.update_request_manager.response():
            if response:
                raise rdb_utils.RDBException('Host %s unable to perform update '
                        '%s through rdb on behalf of %s: %s',  self.hostname,
                        payload, self.dbg_str, response)
        super(RDBClientHostWrapper, self)._update_attributes(payload)
 
 
    def get_metric_fields(self):
        """Generate default set of fields to include for Monarch.
 
        @return: Dictionary of default fields.
        """
        fields = {
            'dut_host_name': self.hostname,
            'board': self.board,
            'model': self._model,
        }
 
        return fields
 
 
    def record_pool(self, fields):
        """Report to Monarch current pool of dut.
 
        @param fields   Dictionary of fields to include.
        """
        pool = ''
        if len(self.pools) == 1:
            pool = self.pools[0]
        if pool in constants.Pools.MANAGED_POOLS:
            pool = 'managed:' + pool
 
        metrics.String(self._HOST_POOL_METRIC,
                       reset_after=True).set(pool, fields=fields)
 
 
    def set_status(self, status):
        """Proxy for setting the status of a host via the rdb.
 
        @param status: The new status.
        """
        # Update elasticsearch db.
        self._update({'status': status})
 
        # Update Monarch.
        fields = self.get_metric_fields()
        self.record_pool(fields)
        # As each device switches state, indicate that it is not in any
        # other state.  This allows Monarch queries to avoid double counting
        # when additional points are added by the Window Align operation.
        host_status_metric = metrics.Boolean(
                self._HOST_STATUS_METRIC, reset_after=True)
        for s in rdb_models.AbstractHostModel.Status.names:
            fields['status'] = s
            host_status_metric.set(s == status, fields=fields)
 
 
    def record_working_state(self, working, timestamp):
        """Report to Monarch whether we are working or broken.
 
        @param working    Host repair status. `True` means that the DUT
                          is up and expected to pass tests.  `False`
                          means the DUT has failed repair and requires
                          manual intervention.
        @param timestamp  Time that the status was recorded.
        """
        fields = self.get_metric_fields()
        metrics.Boolean(
                self._HOST_WORKING_METRIC, reset_after=True).set(
                        working, fields=fields)
        metrics.Boolean(self._BOARD_SHARD_METRIC, reset_after=True).set(
                True,
                fields={
                        'board': self.board,
                        'model': self._model,
                },
        )
        self.record_pool(fields)
 
 
    def update_field(self, fieldname, value):
        """Proxy for updating a field on the host.
 
        @param fieldname: The fieldname as a string.
        @param value: The value to assign to the field.
        """
        self._update({fieldname: value})
 
 
    def platform_and_labels(self):
        """Get the platform and labels on this host.
 
        @return: A tuple containing a list of label names and the platform name.
        """
        platform = self.platform_name
        labels = [label for label in self.labels if label != platform]
        return platform, labels
 
 
    def platform(self):
        """Get the name of the platform of this host.
 
        @return: A string representing the name of the platform.
        """
        return self.platform_name
 
 
    @property
    def board(self):
        """Get the names of the board of this host.
 
        @return: A string of the name of the board, e.g., lumpy. Returns '' if
                no board label is found.
        """
        labels = labellib.LabelsMapping(self.labels)
        return labels.get('board', '')
 
 
    @property
    def _model(self):
        """Get the model this host.
 
        @return: A string of the name of the model, e.g., robo360. Returns '' if
                no model label is found.
        """
        labels = labellib.LabelsMapping(self.labels)
        return labels.get('model', '')
 
 
    @property
    def pools(self):
        """Get the names of the pools of this host.
 
        @return: A list of pool names that the host is assigned to.
        """
        return [l[len(constants.Labels.POOL_PREFIX):] for l in self.labels
                if l.startswith(constants.Labels.POOL_PREFIX)]
 
 
    def get_object_dict(self, **kwargs):
        """Serialize the attributes of this object into a dict.
 
        This method is called through frontend code to get a serialized
        version of this object.
 
        @param kwargs:
            extra_fields: Extra fields, outside the columns of a host table.
 
        @return: A dictionary representing the fields of this host object.
        """
        # TODO(beeps): Implement support for extra fields. Currently nothing
        # requires them.
        return self.wire_format()
 
 
    def save(self):
        """Save any local data a client of this host object might have saved.
 
        Setting attributes on a model before calling its save() method is a
        common django pattern. Most, if not all updates to the host happen
        either through set status or update_field. Though we keep the internal
        state of the RDBClientHostWrapper consistent through these updates
        we need a bulk save method such as this one to save any attributes of
        this host another model might have set on it before calling its own
        save method. Eg:
            task = ST.objects.get(id=12)
            task.host.status = 'Running'
            task.save() -> this should result in the hosts status changing to
            Running.
 
        Functions like add_host_to_labels will have to update this host object
        differently, as that is another level of foreign key indirection.
        """
        self._update(self.get_required_fields_from_host(self))
 
 
def return_rdb_host(func):
    """Decorator for functions that return a list of Host objects.
 
    @param func: The decorated function.
    @return: A functions capable of converting each host_object to a
        rdb_hosts.RDBServerHostWrapper.
    """
    def get_rdb_host(*args, **kwargs):
        """Takes a list of hosts and returns a list of host_infos.
 
        @param hosts: A list of hosts. Each host is assumed to contain
            all the fields in a host_info defined above.
        @return: A list of rdb_hosts.RDBServerHostWrappers, one per host, or an
            empty list is no hosts were found..
        """
        hosts = func(*args, **kwargs)
        return [RDBServerHostWrapper(host) for host in hosts]
    return get_rdb_host