• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python2
2#
3# Copyright 2010 Google Inc. All Rights Reserved.
4"""Script to lock/unlock machines."""
5
6from __future__ import print_function
7
8__author__ = 'asharif@google.com (Ahmad Sharif)'
9
10import argparse
11import datetime
12import fcntl
13import getpass
14import glob
15import json
16import os
17import socket
18import sys
19import time
20
21from cros_utils import logger
22
23LOCK_SUFFIX = '_check_lock_liveness'
24
25# The locks file directory REQUIRES that 'group' only has read/write
26# privileges and 'world' has no privileges.  So the mask must be
27# '0027': 0777 - 0027 = 0750.
28LOCK_MASK = 0027
29
30
31def FileCheckName(name):
32  return name + LOCK_SUFFIX
33
34
35def OpenLiveCheck(file_name):
36  with FileCreationMask(LOCK_MASK):
37    fd = open(file_name, 'a')
38  try:
39    fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
40  except IOError:
41    raise
42  return fd
43
44
45class FileCreationMask(object):
46  """Class for the file creation mask."""
47
48  def __init__(self, mask):
49    self._mask = mask
50    self._old_mask = None
51
52  def __enter__(self):
53    self._old_mask = os.umask(self._mask)
54
55  def __exit__(self, typ, value, traceback):
56    os.umask(self._old_mask)
57
58
59class LockDescription(object):
60  """The description of the lock."""
61
62  def __init__(self, desc=None):
63    try:
64      self.owner = desc['owner']
65      self.exclusive = desc['exclusive']
66      self.counter = desc['counter']
67      self.time = desc['time']
68      self.reason = desc['reason']
69      self.auto = desc['auto']
70    except (KeyError, TypeError):
71      self.owner = ''
72      self.exclusive = False
73      self.counter = 0
74      self.time = 0
75      self.reason = ''
76      self.auto = False
77
78  def IsLocked(self):
79    return self.counter or self.exclusive
80
81  def __str__(self):
82    return ' '.join([
83        'Owner: %s' % self.owner, 'Exclusive: %s' % self.exclusive,
84        'Counter: %s' % self.counter, 'Time: %s' % self.time,
85        'Reason: %s' % self.reason, 'Auto: %s' % self.auto
86    ])
87
88
89class FileLock(object):
90  """File lock operation class."""
91  FILE_OPS = []
92
93  def __init__(self, lock_filename):
94    self._filepath = lock_filename
95    lock_dir = os.path.dirname(lock_filename)
96    assert os.path.isdir(lock_dir), ("Locks dir: %s doesn't exist!" % lock_dir)
97    self._file = None
98    self._description = None
99
100  def getDescription(self):
101    return self._description
102
103  def getFilePath(self):
104    return self._filepath
105
106  def setDescription(self, desc):
107    self._description = desc
108
109  @classmethod
110  def AsString(cls, file_locks):
111    stringify_fmt = '%-30s %-15s %-4s %-4s %-15s %-40s %-4s'
112    header = stringify_fmt % ('machine', 'owner', 'excl', 'ctr', 'elapsed',
113                              'reason', 'auto')
114    lock_strings = []
115    for file_lock in file_locks:
116
117      elapsed_time = datetime.timedelta(
118          seconds=int(time.time() - file_lock.getDescription().time))
119      elapsed_time = '%s ago' % elapsed_time
120      lock_strings.append(
121          stringify_fmt %
122          (os.path.basename(file_lock.getFilePath),
123           file_lock.getDescription().owner,
124           file_lock.getDescription().exclusive,
125           file_lock.getDescription().counter, elapsed_time,
126           file_lock.getDescription().reason, file_lock.getDescription().auto))
127    table = '\n'.join(lock_strings)
128    return '\n'.join([header, table])
129
130  @classmethod
131  def ListLock(cls, pattern, locks_dir):
132    if not locks_dir:
133      locks_dir = Machine.LOCKS_DIR
134    full_pattern = os.path.join(locks_dir, pattern)
135    file_locks = []
136    for lock_filename in glob.glob(full_pattern):
137      if LOCK_SUFFIX in lock_filename:
138        continue
139      file_lock = FileLock(lock_filename)
140      with file_lock as lock:
141        if lock.IsLocked():
142          file_locks.append(file_lock)
143    logger.GetLogger().LogOutput('\n%s' % cls.AsString(file_locks))
144
145  def __enter__(self):
146    with FileCreationMask(LOCK_MASK):
147      try:
148        self._file = open(self._filepath, 'a+')
149        self._file.seek(0, os.SEEK_SET)
150
151        if fcntl.flock(self._file.fileno(), fcntl.LOCK_EX) == -1:
152          raise IOError('flock(%s, LOCK_EX) failed!' % self._filepath)
153
154        try:
155          desc = json.load(self._file)
156        except (EOFError, ValueError):
157          desc = None
158        self._description = LockDescription(desc)
159
160        if self._description.exclusive and self._description.auto:
161          locked_byself = False
162          for fd in self.FILE_OPS:
163            if fd.name == FileCheckName(self._filepath):
164              locked_byself = True
165              break
166          if not locked_byself:
167            try:
168              fp = OpenLiveCheck(FileCheckName(self._filepath))
169            except IOError:
170              pass
171            else:
172              self._description = LockDescription()
173              fcntl.lockf(fp, fcntl.LOCK_UN)
174              fp.close()
175        return self._description
176      # Check this differently?
177      except IOError as ex:
178        logger.GetLogger().LogError(ex)
179        return None
180
181  def __exit__(self, typ, value, traceback):
182    self._file.truncate(0)
183    self._file.write(json.dumps(self._description.__dict__, skipkeys=True))
184    self._file.close()
185
186  def __str__(self):
187    return self.AsString([self])
188
189
190class Lock(object):
191  """Lock class"""
192
193  def __init__(self, lock_file, auto=True):
194    self._to_lock = os.path.basename(lock_file)
195    self._lock_file = lock_file
196    self._logger = logger.GetLogger()
197    self._auto = auto
198
199  def NonBlockingLock(self, exclusive, reason=''):
200    with FileLock(self._lock_file) as lock:
201      if lock.exclusive:
202        self._logger.LogError(
203            'Exclusive lock already acquired by %s. Reason: %s' % (lock.owner,
204                                                                   lock.reason))
205        return False
206
207      if exclusive:
208        if lock.counter:
209          self._logger.LogError('Shared lock already acquired')
210          return False
211        lock_file_check = FileCheckName(self._lock_file)
212        fd = OpenLiveCheck(lock_file_check)
213        FileLock.FILE_OPS.append(fd)
214
215        lock.exclusive = True
216        lock.reason = reason
217        lock.owner = getpass.getuser()
218        lock.time = time.time()
219        lock.auto = self._auto
220      else:
221        lock.counter += 1
222    self._logger.LogOutput('Successfully locked: %s' % self._to_lock)
223    return True
224
225  def Unlock(self, exclusive, force=False):
226    with FileLock(self._lock_file) as lock:
227      if not lock.IsLocked():
228        self._logger.LogWarning("Can't unlock unlocked machine!")
229        return True
230
231      if lock.exclusive != exclusive:
232        self._logger.LogError('shared locks must be unlocked with --shared')
233        return False
234
235      if lock.exclusive:
236        if lock.owner != getpass.getuser() and not force:
237          self._logger.LogError("%s can't unlock lock owned by: %s" %
238                                (getpass.getuser(), lock.owner))
239          return False
240        if lock.auto != self._auto:
241          self._logger.LogError("Can't unlock lock with different -a"
242                                ' parameter.')
243          return False
244        lock.exclusive = False
245        lock.reason = ''
246        lock.owner = ''
247
248        if self._auto:
249          del_list = [
250              i for i in FileLock.FILE_OPS
251              if i.name == FileCheckName(self._lock_file)
252          ]
253          for i in del_list:
254            FileLock.FILE_OPS.remove(i)
255          for f in del_list:
256            fcntl.lockf(f, fcntl.LOCK_UN)
257            f.close()
258          del del_list
259          os.remove(FileCheckName(self._lock_file))
260
261      else:
262        lock.counter -= 1
263    return True
264
265
266class Machine(object):
267  """Machine class"""
268
269  LOCKS_DIR = '/google/data/rw/users/mo/mobiletc-prebuild/locks'
270
271  def __init__(self, name, locks_dir=LOCKS_DIR, auto=True):
272    self._name = name
273    self._auto = auto
274    try:
275      self._full_name = socket.gethostbyaddr(name)[0]
276    except socket.error:
277      self._full_name = self._name
278    self._full_name = os.path.join(locks_dir, self._full_name)
279
280  def Lock(self, exclusive=False, reason=''):
281    lock = Lock(self._full_name, self._auto)
282    return lock.NonBlockingLock(exclusive, reason)
283
284  def TryLock(self, timeout=300, exclusive=False, reason=''):
285    locked = False
286    sleep = timeout / 10
287    while True:
288      locked = self.Lock(exclusive, reason)
289      if locked or not timeout >= 0:
290        break
291      print('Lock not acquired for {0}, wait {1} seconds ...'.format(
292          self._name, sleep))
293      time.sleep(sleep)
294      timeout -= sleep
295    return locked
296
297  def Unlock(self, exclusive=False, ignore_ownership=False):
298    lock = Lock(self._full_name, self._auto)
299    return lock.Unlock(exclusive, ignore_ownership)
300
301
302def Main(argv):
303  """The main function."""
304
305  parser = argparse.ArgumentParser()
306  parser.add_argument(
307      '-r', '--reason', dest='reason', default='', help='The lock reason.')
308  parser.add_argument(
309      '-u',
310      '--unlock',
311      dest='unlock',
312      action='store_true',
313      default=False,
314      help='Use this to unlock.')
315  parser.add_argument(
316      '-l',
317      '--list_locks',
318      dest='list_locks',
319      action='store_true',
320      default=False,
321      help='Use this to list locks.')
322  parser.add_argument(
323      '-f',
324      '--ignore_ownership',
325      dest='ignore_ownership',
326      action='store_true',
327      default=False,
328      help="Use this to force unlock on a lock you don't own.")
329  parser.add_argument(
330      '-s',
331      '--shared',
332      dest='shared',
333      action='store_true',
334      default=False,
335      help='Use this for a shared (non-exclusive) lock.')
336  parser.add_argument(
337      '-d',
338      '--dir',
339      dest='locks_dir',
340      action='store',
341      default=Machine.LOCKS_DIR,
342      help='Use this to set different locks_dir')
343  parser.add_argument('args', nargs='*', help='Machine arg.')
344
345  options = parser.parse_args(argv)
346
347  options.locks_dir = os.path.abspath(options.locks_dir)
348  exclusive = not options.shared
349
350  if not options.list_locks and len(options.args) != 2:
351    logger.GetLogger().LogError(
352        'Either --list_locks or a machine arg is needed.')
353    return 1
354
355  if len(options.args) > 1:
356    machine = Machine(options.args[1], options.locks_dir, auto=False)
357  else:
358    machine = None
359
360  if options.list_locks:
361    FileLock.ListLock('*', options.locks_dir)
362    retval = True
363  elif options.unlock:
364    retval = machine.Unlock(exclusive, options.ignore_ownership)
365  else:
366    retval = machine.Lock(exclusive, options.reason)
367
368  if retval:
369    return 0
370  else:
371    return 1
372
373
374if __name__ == '__main__':
375  sys.exit(Main(sys.argv[1:]))
376