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