improve acquire_lock() method

3 files changed, 74 insertions, 19 deletions

diff --git a/proc.py b/proc.py
index bc29caa..a08969e 100644
--- a/proc.py
+++ b/proc.py
@@ -11,6 +11,7 @@ __docformat__ = "restructuredtext en"
 
 import os
 from os.path import exists
+import stat
 
 from logilab.common.tree import Node
 
@@ -32,14 +33,14 @@ VSIZE = 22
 
 class ProcInfo(Node):
     """provide access to process information found in /proc"""
-    
+
     def __init__(self, pid):
-        Node.__init__(self, pid)
-        self.pid = pid
-        proc_exists(pid)
-        self.file = '/proc/%s/stat' % pid
+        self.pid = int(pid)
+        Node.__init__(self, self.pid)
+        proc_exists(self.pid)
+        self.file = '/proc/%s/stat' % self.pid
         self.ppid = int(self.status()[PPID])
-        
+
     def memory_usage(self):
         """return the memory usage of the process in Ko"""
         try :
@@ -63,7 +64,16 @@ class ProcInfo(Node):
         """return the list of fields found in /proc/<pid>/stat"""
         return open(self.file).read().split()
 
-    
+    def name(self):
+        """return the process name found in /proc/<pid>/stat
+        """
+        return self.status()[1].strip('()')
+
+    def age(self):
+        """return the age of the process
+        """
+        return os.stat(self.file)[stat.ST_MTIME]
+
 class ProcInfoLoader:
     """manage process information"""
     
diff --git a/shellutils.py b/shellutils.py
index 6a7509d..4cfe9a3 100644
--- a/shellutils.py
+++ b/shellutils.py
@@ -15,9 +15,11 @@ import sys
 import tempfile
 import time
 import fnmatch
+import errno
 from os.path import exists, isdir, islink, basename, join, walk
 
 from logilab.common import STD_BLACKLIST
+from logilab.common.proc import ProcInfo, NoSuchProcess
 
 
 def chown(path, login=None, group=None):
@@ -199,21 +201,37 @@ class Execute:
         os.remove(outfile)
         os.remove(errfile)
 
-
 def acquire_lock(lock_file, max_try=10, delay=10):
     """Acquire a lock represented by a file on the file system."""
-    count = 0
-    while max_try <= 0 or count < max_try:
-        if not exists(lock_file):
-            break
-        count += 1
-        time.sleep(delay)
+    count = abs(max_try)
+    while count:
+        try:
+            fd = os.open(lock_file, os.O_EXCL | os.O_RDWR | os.O_CREAT)
+            os.write(fd, str(os.getpid()))
+            os.close(fd)
+            return True
+        except OSError, e:
+            if e.errno == errno.EEXIST:
+                try:
+                    fd = open(lock_file, "r")
+                    pid = int(fd.readline())
+                    pi = ProcInfo(pid)
+                    # only print the message one time
+                    if count == max_try:
+                        diff = (int(time.time()) - pi.age()) / 60
+                        print("Command '%s' (pid %s) have locked the file '%s' for %s minutes.\nWaiting..." % (pi.name(), pid, lock_file, diff))
+                except NoSuchProcess:
+                    raise NoSuchProcess('You can delete the lock file "%s" safely' % lock_file)
+                except:
+                    # process information are not accessible
+                    pass
+            else:
+                raise
+            count -= 1
+            time.sleep(delay)
     else:
         raise Exception('Unable to acquire %s' % lock_file)
-    stream = open(lock_file, 'w')
-    stream.write(str(os.getpid()))
-    stream.close()
-    
+
 def release_lock(lock_file):
     """Release a lock represented by a file on the file system."""
     os.remove(lock_file)
diff --git a/test/unittest_shellutils.py b/test/unittest_shellutils.py
index 396d7c6..903312d 100644
--- a/test/unittest_shellutils.py
+++ b/test/unittest_shellutils.py
@@ -5,7 +5,8 @@ from os.path import join
 
 from logilab.common.testlib import TestCase, unittest_main
 
-from logilab.common.shellutils import globfind, find, ProgressBar
+from logilab.common.shellutils import globfind, find, ProgressBar, acquire_lock, release_lock
+from logilab.common.proc import NoSuchProcess
 from StringIO import StringIO
 
 DATA_DIR = join('data','find_test')
@@ -123,5 +124,31 @@ class ProgressBarTC(TestCase):
         self._update_test(5, (8, 16, 25, 33, 42, (42, True)), size=42)
 
 
+class AcquireLockTC(TestCase):
+
+    def setUp(self):
+        self.tmpdir = tempfile.mkdtemp()
+        self.lock = join(self.tmpdir, 'LOCK')
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpdir)
+
+    def test_acquire_normal(self):
+        self.assertTrue(acquire_lock(self.lock, 1, 1))
+        self.assertTrue(os.path.exists(self.lock))
+        release_lock(self.lock)
+        self.assertFalse(os.path.exists(self.lock))
+
+    def test_no_possible_acquire(self):
+        self.assertRaises(Exception, acquire_lock, self.lock, 0)
+
+    def test_wrong_process(self):
+        fd = os.open(self.lock, os.O_EXCL | os.O_RDWR | os.O_CREAT)
+        os.write(fd, '1111111111')
+        os.close(fd)
+        self.assertTrue(os.path.exists(self.lock))
+        self.assertRaises(NoSuchProcess, acquire_lock, self.lock)
+
+
 if __name__ == '__main__':
     unittest_main()