file_lock_machine.py revision 88272d479f2761cc1906fea564c73033f77a6270 
1#!/usr/bin/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(['Owner: %s' % self.owner, 'Exclusive: %s' % self.exclusive,
83                     'Counter: %s' % self.counter, 'Time: %s' % self.time,
84                     'Reason: %s' % self.reason, 'Auto: %s' % self.auto])
85
86
87class FileLock(object):
88  """File lock operation class."""
89  FILE_OPS = []
90
91  def __init__(self, lock_filename):
92    self._filepath = lock_filename
93    lock_dir = os.path.dirname(lock_filename)
94    assert os.path.isdir(lock_dir), ("Locks dir: %s doesn't exist!" % lock_dir)
95    self._file = None
96    self._description = None
97
98  def getDescription(self):
99    return self._description
100
101  def getFilePath(self):
102    return self._filepath
103
104  def setDescription(self, desc):
105    self._description = desc
106
107  @classmethod
108  def AsString(cls, file_locks):
109    stringify_fmt = '%-30s %-15s %-4s %-4s %-15s %-40s %-4s'
110    header = stringify_fmt % ('machine', 'owner', 'excl', 'ctr', 'elapsed',
111                              'reason', 'auto')
112    lock_strings = []
113    for file_lock in file_locks:
114
115      elapsed_time = datetime.timedelta(
116          seconds=int(time.time() - file_lock.getDescription().time))
117      elapsed_time = '%s ago' % elapsed_time
118      lock_strings.append(
119          stringify_fmt %
120          (os.path.basename(file_lock.getFilePath),
121           file_lock.getDescription().owner,
122           file_lock.getDescription().exclusive,
123           file_lock.getDescription().counter,
124           elapsed_time, file_lock.getDescription().reason,
125           file_lock.getDescription().auto))
126    table = '\n'.join(lock_strings)
127    return '\n'.join([header, table])
128
129  @classmethod
130  def ListLock(cls, pattern, locks_dir):
131    if not locks_dir:
132      locks_dir = Machine.LOCKS_DIR
133    full_pattern = os.path.join(locks_dir, pattern)
134    file_locks = []
135    for lock_filename in glob.glob(full_pattern):
136      if LOCK_SUFFIX in lock_filename:
137        continue
138      file_lock = FileLock(lock_filename)
139      with file_lock as lock:
140        if lock.IsLocked():
141          file_locks.append(file_lock)
142    logger.GetLogger().LogOutput('\n%s' % cls.AsString(file_locks))
143
144  def __enter__(self):
145    with FileCreationMask(LOCK_MASK):
146      try:
147        self._file = open(self._filepath, 'a+')
148        self._file.seek(0, os.SEEK_SET)
149
150        if fcntl.flock(self._file.fileno(), fcntl.LOCK_EX) == -1:
151          raise IOError('flock(%s, LOCK_EX) failed!' % self._filepath)
152
153        try:
154          desc = json.load(self._file)
155        except (EOFError, ValueError):
156          desc = None
157        self._description = LockDescription(desc)
158
159        if self._description.exclusive and self._description.auto:
160          locked_byself = False
161          for fd in self.FILE_OPS:
162            if fd.name == FileCheckName(self._filepath):
163              locked_byself = True
164              break
165          if not locked_byself:
166            try:
167              fp = OpenLiveCheck(FileCheckName(self._filepath))
168            except IOError:
169              pass
170            else:
171              self._description = LockDescription()
172              fcntl.lockf(fp, fcntl.LOCK_UN)
173              fp.close()
174        return self._description
175      # Check this differently?
176      except IOError as ex:
177        logger.GetLogger().LogError(ex)
178        return None
179
180  def __exit__(self, typ, value, traceback):
181    self._file.truncate(0)
182    self._file.write(json.dumps(self._description.__dict__, skipkeys=True))
183    self._file.close()
184
185  def __str__(self):
186    return self.AsString([self])
187
188
189class Lock(object):
190  """Lock class"""
191
192  def __init__(self, lock_file, auto=True):
193    self._to_lock = os.path.basename(lock_file)
194    self._lock_file = lock_file
195    self._logger = logger.GetLogger()
196    self._auto = auto
197
198  def NonBlockingLock(self, exclusive, reason=''):
199    with FileLock(self._lock_file) as lock:
200      if lock.exclusive:
201        self._logger.LogError(
202            'Exclusive lock already acquired by %s. Reason: %s' %
203            (lock.owner, lock.reason))
204        return False
205
206      if exclusive:
207        if lock.counter:
208          self._logger.LogError('Shared lock already acquired')
209          return False
210        lock_file_check = FileCheckName(self._lock_file)
211        fd = OpenLiveCheck(lock_file_check)
212        FileLock.FILE_OPS.append(fd)
213
214        lock.exclusive = True
215        lock.reason = reason
216        lock.owner = getpass.getuser()
217        lock.time = time.time()
218        lock.auto = self._auto
219      else:
220        lock.counter += 1
221    self._logger.LogOutput('Successfully locked: %s' % self._to_lock)
222    return True
223
224  def Unlock(self, exclusive, force=False):
225    with FileLock(self._lock_file) as lock:
226      if not lock.IsLocked():
227        self._logger.LogWarning("Can't unlock unlocked machine!")
228        return True
229
230      if lock.exclusive != exclusive:
231        self._logger.LogError('shared locks must be unlocked with --shared')
232        return False
233
234      if lock.exclusive:
235        if lock.owner != getpass.getuser() and not force:
236          self._logger.LogError("%s can't unlock lock owned by: %s" %
237                                (getpass.getuser(), lock.owner))
238          return False
239        if lock.auto != self._auto:
240          self._logger.LogError("Can't unlock lock with different -a"
241                                ' parameter.')
242          return False
243        lock.exclusive = False
244        lock.reason = ''
245        lock.owner = ''
246
247        if self._auto:
248          del_list = [i
249                      for i in FileLock.FILE_OPS
250                      if i.name == FileCheckName(self._lock_file)]
251          for i in del_list:
252            FileLock.FILE_OPS.remove(i)
253          for f in del_list:
254            fcntl.lockf(f, fcntl.LOCK_UN)
255            f.close()
256          del del_list
257          os.remove(FileCheckName(self._lock_file))
258
259      else:
260        lock.counter -= 1
261    return True
262
263
264class Machine(object):
265  """Machine class"""
266
267  LOCKS_DIR = '/google/data/rw/users/mo/mobiletc-prebuild/locks'
268
269  def __init__(self, name, locks_dir=LOCKS_DIR, auto=True):
270    self._name = name
271    self._auto = auto
272    try:
273      self._full_name = socket.gethostbyaddr(name)[0]
274    except socket.error:
275      self._full_name = self._name
276    self._full_name = os.path.join(locks_dir, self._full_name)
277
278  def Lock(self, exclusive=False, reason=''):
279    lock = Lock(self._full_name, self._auto)
280    return lock.NonBlockingLock(exclusive, reason)
281
282  def TryLock(self, timeout=300, exclusive=False, reason=''):
283    locked = False
284    sleep = timeout / 10
285    while True:
286      locked = self.Lock(exclusive, reason)
287      if locked or not timeout >= 0:
288        break
289      print('Lock not acquired for {0}, wait {1} seconds ...'.format(
290          self._name,
291          sleep))
292      time.sleep(sleep)
293      timeout -= sleep
294    return locked
295
296  def Unlock(self, exclusive=False, ignore_ownership=False):
297    lock = Lock(self._full_name, self._auto)
298    return lock.Unlock(exclusive, ignore_ownership)
299
300
301def Main(argv):
302  """The main function."""
303
304  parser = argparse.ArgumentParser()
305  parser.add_argument('-r',
306                      '--reason',
307                      dest='reason',
308                      default='',
309                      help='The lock reason.')
310  parser.add_argument('-u',
311                      '--unlock',
312                      dest='unlock',
313                      action='store_true',
314                      default=False,
315                      help='Use this to unlock.')
316  parser.add_argument('-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('-f',
323                      '--ignore_ownership',
324                      dest='ignore_ownership',
325                      action='store_true',
326                      default=False,
327                      help="Use this to force unlock on a lock you don't own.")
328  parser.add_argument('-s',
329                      '--shared',
330                      dest='shared',
331                      action='store_true',
332                      default=False,
333                      help='Use this for a shared (non-exclusive) lock.')
334  parser.add_argument('-d',
335                      '--dir',
336                      dest='locks_dir',
337                      action='store',
338                      default=Machine.LOCKS_DIR,
339                      help='Use this to set different locks_dir')
340  parser.add_argument('args', nargs='*', help='Machine arg.')
341
342  options = parser.parse_args(argv)
343
344  options.locks_dir = os.path.abspath(options.locks_dir)
345  exclusive = not options.shared
346
347  if not options.list_locks and len(options.args) != 2:
348    logger.GetLogger().LogError(
349        'Either --list_locks or a machine arg is needed.')
350    return 1
351
352  if len(options.args) > 1:
353    machine = Machine(options.args[1], options.locks_dir, auto=False)
354  else:
355    machine = None
356
357  if options.list_locks:
358    FileLock.ListLock('*', options.locks_dir)
359    retval = True
360  elif options.unlock:
361    retval = machine.Unlock(exclusive, options.ignore_ownership)
362  else:
363    retval = machine.Lock(exclusive, options.reason)
364
365  if retval:
366    return 0
367  else:
368    return 1
369
370
371if __name__ == '__main__':
372  sys.exit(Main(sys.argv[1:]))
373