1
0

dlock.py 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219
  1. #!/usr/bin/python
  2. #
  3. # Copyright 2006-2007 BitTorrent, Inc. All Rights Reserved.
  4. #
  5. # Written by Ben Teitelbaum
  6. import os
  7. import sys
  8. import socket
  9. from time import asctime, gmtime, time, sleep
  10. from twisted.internet import reactor, task
  11. class dlock(object):
  12. def __init__(self, deadlockfile, update_period=300, myhost=None, debug=None):
  13. if myhost == None: myhost = socket.gethostname()
  14. self.host = myhost
  15. self.pid = os.getpid()
  16. self.deadlockfile = deadlockfile
  17. self.refresher = task.LoopingCall(self.refresh)
  18. self.update_period = update_period
  19. self.debug = debug
  20. # Block until lock is acquired, then refresh the lock file every
  21. # update_period seconds.
  22. #
  23. # Nota Bene: while blocked on acquiring the lock, this sleeps the
  24. # whole process; once the lock is acquired, an event-driven model
  25. # (twisted reactor) is presumed. The intended use (see test at
  26. # bottom) is to block on acquire before running the Twisted
  27. # reactor.
  28. #
  29. def acquire(self):
  30. while True:
  31. while self.islocked():
  32. if self.debug:
  33. lock = self._readlock()
  34. print '%s locked by %s' % (self.deadlockfile, self._lockdict2string(lock))
  35. sleep(self.update_period)
  36. try:
  37. # Use link count hack to work around NFS's broken
  38. # file locking.
  39. tempfile = '.' + str(self.pid) + self.host + str(time()) + '.tmp'
  40. lockfile = self.deadlockfile + '.lock'
  41. # Create temp lock file
  42. fh = open(tempfile, "w")
  43. fh.close()
  44. # Atomicallly create lockfile as a hard link
  45. try:
  46. os.link(tempfile, lockfile)
  47. except:
  48. if self.debug:
  49. print "tempfile: " + tempfile
  50. print "lockfile: " + lockfile
  51. raise
  52. # Check the number of links
  53. if os.stat(tempfile)[3] == os.stat(lockfile)[3]:
  54. # Hooray, I have the write lock on the deadlock file!
  55. self._timestamp_deadlockfile(time())
  56. if self.debug:
  57. lock = self._readlock()
  58. print '%s acquired by %s' % (self.deadlockfile, self._lockdict2string(lock))
  59. self.refresher.start(self.update_period)
  60. # Release the lock
  61. os.unlink(tempfile)
  62. os.unlink(lockfile)
  63. return self
  64. else:
  65. # Failed to grab write lock on deadlock file, keep looping
  66. if self.debug:
  67. print '%d failed to grab write lock on deadlock file: %s (will retry)' % (self.pid, self.deadlockfile)
  68. except:
  69. if self.debug:
  70. print 'File Lock Error: %s@%s could not acquire %s' % (self.pid, self.host, self.deadlockfile)
  71. raise
  72. def refresh(self):
  73. assert self.ownlock()
  74. # No need to grab a write lock on the deadlock file, since it's not stale
  75. self._timestamp_deadlockfile(time())
  76. def _timestamp_deadlockfile(self, ts):
  77. try:
  78. fh = open(self.deadlockfile, 'w')
  79. fh.write(self._lockstr(ts))
  80. fh.close()
  81. os.chmod(self.deadlockfile, 0644)
  82. except:
  83. if self.debug:
  84. print 'File Lock Error: %s@%s could not write %s' % (self.pid, self.host, self.deadlockfile)
  85. raise
  86. def release(self):
  87. if self.ownlock():
  88. try:
  89. self.refresher.stop()
  90. self._timestamp_deadlockfile(0)
  91. if self.debug:
  92. print '%s@%s released lock %s' % (self.pid, self.host, self.deadlockfile)
  93. except:
  94. if self.debug:
  95. print 'File Lock Error: %s@%s could not release %s' % (self.pid, self.host, self.deadlockfile)
  96. raise
  97. return self
  98. def islocked(self):
  99. try:
  100. if self._isstale():
  101. # Lock seems stale, wait for one more update period and check again
  102. sleep(self.update_period)
  103. return not self._isstale()
  104. else:
  105. return True
  106. except:
  107. if self.debug:
  108. print "islocked exception"
  109. return False
  110. def _isstale(self):
  111. lock = self._readlock()
  112. if time() - lock['timestamp'] > self.update_period:
  113. return True
  114. else:
  115. return False
  116. def _readlock(self):
  117. try:
  118. lock = {}
  119. fh = open(self.deadlockfile)
  120. data = fh.read().split()
  121. fh.close()
  122. assert len(data) == 3
  123. lock['pid'] = int(data[0])
  124. lock['host'] = data[1]
  125. lock['timestamp'] = float(data[2])
  126. return lock
  127. except:
  128. if self.debug:
  129. print 'File Lock Error: %s@%s reading %s' % (self.pid, self.host, self.deadlockfile)
  130. raise
  131. # Public method to read a lockfile.
  132. @classmethod
  133. def readlock(cls, lockfile):
  134. lock = cls(deadlockfile=lockfile, myhost='dummy')
  135. return lock._readlock()
  136. def _lockdict2string(self, lock):
  137. return '%s@%s at %s' % (lock['pid'], lock['host'], asctime(gmtime(lock['timestamp'])))
  138. def _lockstr(self, ts):
  139. return '%d %s %f'%(self.pid, self.host, ts)
  140. def ownlock(self):
  141. lock = self._readlock()
  142. return (self.host == lock['host'] and
  143. self.pid == lock['pid'])
  144. def __del__(self):
  145. self.release()
  146. # Tests
  147. #
  148. # Run several in parallel on multiple machines, but have at most one
  149. # whack the deadlock file on initialization.
  150. #
  151. def run_tests(argv=None):
  152. if argv is None:
  153. argv = sys.argv
  154. deadlockfile = './dlock_test'
  155. l = dlock(deadlockfile, 5, debug=True)
  156. # Stupid argv handling; just grab first arg and run that test
  157. if len(argv) > 1:
  158. if argv[1] == 'none':
  159. print "Removing deadlock file."
  160. os.unlink(deadlockfile)
  161. elif argv[1] == 'old':
  162. print "Creating stale deadlock file owned by no one."
  163. fh = open(l.deadlockfile, 'w')
  164. fh.write('%d %s %f'%(0, 0, 0))
  165. fh.close()
  166. elif argv[1] == 'new':
  167. print "Creating fresh deadlock file owned by no one."
  168. fh = open(l.deadlockfile, 'w')
  169. fh.write('%d %s %f'%(0, 0, time()))
  170. fh.close()
  171. else:
  172. print "Un-known arg--starting with old deadlock file."
  173. else:
  174. print "Starting with old deadlock file."
  175. # Tease for a while, then release the lock
  176. def tease(l, n):
  177. if n > 0:
  178. assert l.ownlock()
  179. print 'I (%d) have the lock--ha, ha ha!'%os.getpid()
  180. reactor.callLater(1, tease, l, n - 1)
  181. else:
  182. l.release()
  183. # Start teasing once reactor is run
  184. reactor.callLater(1, tease, l, 20)
  185. # But first, grab the lock (this blocks)
  186. l.acquire()
  187. reactor.run()
  188. if __name__ == "__main__":
  189. sys.exit(run_tests())