• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
5import logging
6import signal
7import common
8
9from autotest_lib.server import site_utils
10from autotest_lib.server.cros.chaos_lib import chaos_datastore_utils
11"""HostLockManager class, for the dynamic_suite module.
12
13A HostLockManager instance manages locking and unlocking a set of autotest DUTs.
14A caller can lock or unlock one or more DUTs. If the caller fails to unlock()
15locked hosts before the instance is destroyed, it will attempt to unlock() the
16hosts automatically, but this is to be avoided.
17
18Sample usage:
19  manager = host_lock_manager.HostLockManager()
20  try:
21      manager.lock(['host1'])
22      # do things
23  finally:
24      manager.unlock()
25"""
26
27class HostLockManager(object):
28    """
29    @attribute _afe: an instance of AFE as defined in server/frontend.py.
30    @attribute _locked_hosts: a set of DUT hostnames.
31    @attribute LOCK: a string.
32    @attribute UNLOCK: a string.
33    """
34
35    LOCK = 'lock'
36    UNLOCK = 'unlock'
37
38
39    @property
40    def locked_hosts(self):
41        """@returns set of locked hosts."""
42        return self._locked_hosts
43
44
45    @locked_hosts.setter
46    def locked_hosts(self, hosts):
47        """Sets value of locked_hosts.
48
49        @param hosts: a set of strings.
50        """
51        self._locked_hosts = hosts
52
53
54    def __init__(self, afe=None):
55        """
56        Constructor
57        """
58        self.dutils = chaos_datastore_utils.ChaosDataStoreUtils()
59        # Keep track of hosts locked by this instance.
60        self._locked_hosts = set()
61
62
63    def __del__(self):
64        if self._locked_hosts:
65            logging.warning('Caller failed to unlock %r! Forcing unlock now.',
66                            self._locked_hosts)
67            self.unlock()
68
69
70    def _check_host(self, host, operation):
71        """Checks host for desired operation.
72
73        @param host: a string, hostname.
74        @param operation: a string, LOCK or UNLOCK.
75        @returns a string: host name, if desired operation can be performed on
76                           host or None otherwise.
77        """
78        host_checked = host
79        # Get host details from DataStore
80        host_info = self.dutils.show_device(host)
81
82        if not host_info:
83            logging.warning('Host (AP) details not found in DataStore')
84            return None
85
86        if operation == self.LOCK and host_info['lock_status']:
87            err = ('Contention detected: %s is locked by %s at %s.' %
88                   (host, host_info['locked_by'],
89                    host_info['lock_status_updated']))
90            logging.error(err)
91            return None
92
93        elif operation == self.UNLOCK and not host_info['lock_status']:
94            logging.info('%s not locked.', host)
95            return None
96
97        return host_checked
98
99
100    def lock(self, hosts, lock_reason='Locked by HostLockManager'):
101        """Lock hosts in datastore.
102
103        @param hosts: a list of strings, host names.
104        @param lock_reason: a string, a reason for locking the hosts.
105
106        @returns a boolean, True == at least one host from hosts is locked.
107        """
108        # Filter out hosts that we may have already locked
109        new_hosts = set(hosts).difference(self._locked_hosts)
110        logging.info('Attempt to lock %s', new_hosts)
111        if not new_hosts:
112            return False
113
114        return self._host_modifier(new_hosts, self.LOCK, lock_reason=lock_reason)
115
116
117    def unlock(self, hosts=None):
118        """Unlock hosts in datastore after use.
119
120        @param hosts: a list of strings, host names.
121        @returns a boolean, True == at least one host from self._locked_hosts is
122                 unlocked.
123        """
124        # Filter out hosts that we did not lock
125        updated_hosts = self._locked_hosts
126        if hosts:
127            unknown_hosts = set(hosts).difference(self._locked_hosts)
128            logging.warning('Skip unknown hosts: %s', unknown_hosts)
129            updated_hosts = set(hosts) - unknown_hosts
130            logging.info('Valid hosts: %s', updated_hosts)
131            updated_hosts = updated_hosts.intersection(self._locked_hosts)
132
133        if not updated_hosts:
134            return False
135
136        logging.info('Unlocking hosts (APs / PCAPs): %s', updated_hosts)
137        return self._host_modifier(updated_hosts, self.UNLOCK)
138
139
140    def _host_modifier(self, hosts, operation, lock_reason=None):
141        """Helper that locks hosts in DataStore.
142
143        @param: hosts, a set of strings, host names.
144        @param operation: a string, LOCK or UNLOCK.
145        @param lock_reason: a string, a reason must be provided when locking.
146
147        @returns a boolean, if operation succeeded on at least one host in
148                 hosts.
149        """
150        updated_hosts = set()
151        for host in hosts:
152            verified_host = self._check_host(host, operation)
153            if verified_host is not None:
154                updated_hosts.add(verified_host)
155
156        logging.info('host_modifier: updated_hosts = %s', updated_hosts)
157        if not updated_hosts:
158            logging.info('host_modifier: no host to update')
159            return False
160
161        for host in updated_hosts:
162            if operation == self.LOCK:
163                if self.dutils.lock_device(host, lock_reason):
164                    logging.info('Locked host in datastore: %s', host)
165                    self._locked_hosts = self._locked_hosts.union([host])
166                else:
167                    logging.error('Unable to lock host: ', host)
168
169            if operation == self.UNLOCK:
170                if self.dutils.unlock_device(host):
171                    logging.info('Unlocked host in datastore: %s', host)
172                    self._locked_hosts = self._locked_hosts.difference([host])
173                else:
174                    logging.error('Unable to un-lock host: %s', host)
175
176        return True
177
178
179class HostsLockedBy(object):
180    """Context manager to make sure that a HostLockManager will always unlock
181    its machines. This protects against both exceptions and SIGTERM."""
182
183    def _make_handler(self):
184        def _chaining_signal_handler(signal_number, frame):
185            self._manager.unlock()
186            # self._old_handler can also be signal.SIG_{IGN,DFL} which are ints.
187            if callable(self._old_handler):
188                self._old_handler(signal_number, frame)
189        return _chaining_signal_handler
190
191
192    def __init__(self, manager):
193        """
194        @param manager: The HostLockManager used to lock the hosts.
195        """
196        self._manager = manager
197        self._old_handler = signal.SIG_DFL
198
199
200    def __enter__(self):
201        self._old_handler = signal.signal(signal.SIGTERM, self._make_handler())
202
203
204    def __exit__(self, exntype, exnvalue, backtrace):
205        signal.signal(signal.SIGTERM, self._old_handler)
206        self._manager.unlock()
207