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