blob: c5f980927d96aa767d980e548be2cdf63efc09f4 [file] [log] [blame]
Ahmad Sharif0dcbc4b2012-02-02 16:37:18 -08001#!/usr/bin/python2.6
2#
3# Copyright 2010 Google Inc. All Rights Reserved.
4
5"""Script to lock/unlock machines.
6
7"""
8
9__author__ = "asharif@google.com (Ahmad Sharif)"
10
11import datetime
12import fcntl
13import getpass
14import glob
15import optparse
16import os
17import pickle
18import socket
19import sys
20import time
21from utils import logger
22
23
24class FileCreationMask(object):
25 def __init__(self, mask):
26 self._mask = mask
27
28 def __enter__(self):
29 self._old_mask = os.umask(self._mask)
30
31 def __exit__(self, type, value, traceback):
32 os.umask(self._old_mask)
33
34
35class LockDescription(object):
36 def __init__(self):
37 self.owner = ""
38 self.exclusive = False
39 self.counter = 0
40 self.time = 0
41 self.reason = ""
42
43 def IsLocked(self):
44 return self.counter or self.exclusive
45
46 def __str__(self):
47 return " ".join(["Owner: %s" % self.owner,
48 "Exclusive: %s" % self.exclusive,
49 "Counter: %s" % self.counter,
50 "Time: %s" % self.time,
51 "Reason: %s" % self.reason])
52
53
54class FileLock(object):
55 LOCKS_DIR = "/home/mobiletc-prebuild/locks"
56
57 def __init__(self, lock_filename):
58 assert os.path.isdir(self.LOCKS_DIR), (
59 "Locks dir: %s doesn't exist!" % self.LOCKS_DIR)
60 self._filepath = os.path.join(self.LOCKS_DIR, lock_filename)
61 self._file = None
62
63 @classmethod
64 def AsString(cls, file_locks):
65 stringify_fmt = "%-30s %-15s %-4s %-4s %-15s %-40s"
66 header = stringify_fmt % ("machine", "owner", "excl", "ctr",
67 "elapsed", "reason")
68 lock_strings = []
69 for file_lock in file_locks:
70
71 elapsed_time = datetime.timedelta(
72 seconds=int(time.time() - file_lock._description.time))
73 elapsed_time = "%s ago" % elapsed_time
74 lock_strings.append(stringify_fmt %
75 (os.path.basename(file_lock._filepath),
76 file_lock._description.owner,
77 file_lock._description.exclusive,
78 file_lock._description.counter,
79 elapsed_time,
80 file_lock._description.reason))
81 table = "\n".join(lock_strings)
82 return "\n".join([header, table])
83
84 @classmethod
85 def ListLock(cls, pattern):
86 full_pattern = os.path.join(cls.LOCKS_DIR, pattern)
87 file_locks = []
88 for lock_filename in glob.glob(full_pattern):
89 file_lock = FileLock(lock_filename)
90 with file_lock as lock:
91 if lock.IsLocked():
92 file_locks.append(file_lock)
93 logger.GetLogger().LogOutput("\n%s" % cls.AsString(file_locks))
94
95 def __enter__(self):
96 with FileCreationMask(0000):
97 try:
98 self._file = open(self._filepath, "a+")
99 self._file.seek(0, os.SEEK_SET)
100
101 if fcntl.flock(self._file.fileno(), fcntl.LOCK_EX) == -1:
102 raise IOError("flock(%s, LOCK_EX) failed!" % self._filepath)
103
104 try:
105 self._description = pickle.load(self._file)
106 except (EOFError, pickle.PickleError):
107 self._description = LockDescription()
108 return self._description
109 # Check this differently?
110 except IOError as ex:
111 logger.GetLogger().LogError(ex)
112 return None
113
114 def __exit__(self, type, value, traceback):
115 self._file.truncate(0)
116 self._file.write(pickle.dumps(self._description))
117 self._file.close()
118
119 def __str__(self):
120 return self.AsString([self])
121
122
123class Lock(object):
124 def __init__(self, to_lock):
125 self._to_lock = to_lock
126 self._logger = logger.GetLogger()
127
128 def NonBlockingLock(self, exclusive, reason=""):
129 with FileLock(self._to_lock) as lock:
130 if lock.exclusive:
131 self._logger.LogError(
132 "Exclusive lock already acquired by %s. Reason: %s" %
133 (lock.owner, lock.reason))
134 return False
135
136 if exclusive:
137 if lock.counter:
138 self._logger.LogError("Shared lock already acquired")
139 return False
140 lock.exclusive = True
141 lock.reason = reason
142 lock.owner = getpass.getuser()
143 lock.time = time.time()
144 else:
145 lock.counter += 1
146 self._logger.LogOutput("Successfully locked: %s" % self._to_lock)
147 return True
148
149 def Unlock(self, exclusive, force=False):
150 with FileLock(self._to_lock) as lock:
151 if not lock.IsLocked():
152 self._logger.LogError("Can't unlock unlocked machine!")
153 return False
154
155 if lock.exclusive != exclusive:
156 self._logger.LogError("shared locks must be unlocked with --shared")
157 return False
158
159 if lock.exclusive:
160 if lock.owner != getpass.getuser() and not force:
161 self._logger.LogError("%s can't unlock lock owned by: %s" %
162 (getpass.getuser(), lock.owner))
163 return False
164 lock.exclusive = False
165 lock.reason = ""
166 lock.owner = ""
167 else:
168 lock.counter -= 1
169 return True
170
171
172class Machine(object):
173 def __init__(self, name):
174 self._name = name
175 try:
176 self._full_name = socket.gethostbyaddr(name)[0]
177 except socket.error:
178 self._full_name = self._name
179
180 def Lock(self, exclusive=False, reason=""):
181 lock = Lock(self._full_name)
182 return lock.NonBlockingLock(exclusive, reason)
183
184 def Unlock(self, exclusive=False, ignore_ownership=False):
185 lock = Lock(self._full_name)
186 return lock.Unlock(exclusive, ignore_ownership)
187
188
189def Main(argv):
190 """The main function."""
191 parser = optparse.OptionParser()
192 parser.add_option("-r",
193 "--reason",
194 dest="reason",
195 default="",
196 help="The lock reason.")
197 parser.add_option("-u",
198 "--unlock",
199 dest="unlock",
200 action="store_true",
201 default=False,
202 help="Use this to unlock.")
203 parser.add_option("-l",
204 "--list_locks",
205 dest="list_locks",
206 action="store_true",
207 default=False,
208 help="Use this to list locks.")
209 parser.add_option("-f",
210 "--ignore_ownership",
211 dest="ignore_ownership",
212 action="store_true",
213 default=False,
214 help="Use this to force unlock on a lock you don't own.")
215 parser.add_option("-s",
216 "--shared",
217 dest="shared",
218 action="store_true",
219 default=False,
220 help="Use this for a shared (non-exclusive) lock.")
221
222 options, args = parser.parse_args(argv)
223
224 exclusive = not options.shared
225
226 if not options.list_locks and len(args) != 2:
227 logger.GetLogger().LogError(
228 "Either --list_locks or a machine arg is needed.")
229 return 1
230
231 if len(args) > 1:
232 machine = Machine(args[1])
233 else:
234 machine = None
235
236 if options.list_locks:
237 FileLock.ListLock("*")
238 retval = True
239 elif options.unlock:
240 retval = machine.Unlock(exclusive, options.ignore_ownership)
241 else:
242 retval = machine.Lock(exclusive, options.reason)
243
244 if retval:
245 return 0
246 else:
247 return 1
248
249if __name__ == "__main__":
250 sys.exit(Main(sys.argv))