lockfile.py

Go to the documentation of this file.

##
# 
# lockfile.py - Platform-independent advisory file locks.
# 
# Requires Python 2.5 unless you apply 2.4.diff
# Locking is done on a per-thread basis instead of a per-process basis.
# 
# Usage:
# 
# >>> lock = FileLock(_testfile())
# >>> try:
# ...     lock.acquire()
# ... except AlreadyLocked:
# ...     print _testfile(), 'is locked already.'
# ... except LockFailed:
# ...     print _testfile(), 'can\\'t be locked.'
# ... else:
# ...     print 'got lock'
# got lock
# >>> print lock.is_locked()
# True
# >>> lock.release()
# 
# >>> lock = FileLock(_testfile())
# >>> print lock.is_locked()
# False
# >>> with lock:
# ...    print lock.is_locked()
# True
# >>> print lock.is_locked()
# False
# >>> # It is okay to lock twice from the same thread...
# >>> with lock:
# ...     lock.acquire()
# ...
# >>> # Though no counter is kept, so you can't unlock multiple times...
# >>> print lock.is_locked()
# False
# 
# Exceptions:
# 
#     Error - base class for other exceptions
#         LockError - base class for all locking exceptions
#             AlreadyLocked - Another thread or process already holds the lock
#             LockFailed - Lock failed for some other reason
#         UnlockError - base class for all unlocking exceptions
#             AlreadyUnlocked - File was not locked.
#             NotMyLock - File was locked but not by the current thread/process
# 
# To do:
#     * Write more test cases
#       - verify that all lines of code are executed
#     * Describe on-disk file structures in the documentation.
# 

from __future__ import division, with_statement

import socket
import os
import threading
import time
import errno
import thread

##
# 
#     Base class for other exceptions.
# 
#     >>> try:
#     ...   raise Error
#     ... except Exception:
#     ...   pass
#     
class Error(Exception):
    pass

##
# 
#     Base class for error arising from attempts to acquire the lock.
# 
#     >>> try:
#     ...   raise LockError
#     ... except Error:
#     ...   pass
#     
class LockError(Error):
    pass

##
# Raised when lock creation fails within a user-defined period of time.
# 
#     >>> try:
#     ...   raise LockTimeout
#     ... except LockError:
#     ...   pass
#     
class LockTimeout(LockError):
    pass

##
# Some other thread/process is locking the file.
# 
#     >>> try:
#     ...   raise AlreadyLocked
#     ... except LockError:
#     ...   pass
#     
class AlreadyLocked(LockError):
    pass

##
# Lock file creation failed for some other reason.
# 
#     >>> try:
#     ...   raise LockFailed
#     ... except LockError:
#     ...   pass
#     
class LockFailed(LockError):
    pass

##
# 
#     Base class for errors arising from attempts to release the lock.
# 
#     >>> try:
#     ...   raise UnlockError
#     ... except Error:
#     ...   pass
#     
class UnlockError(Error):
    pass

##
# Raised when an attempt is made to unlock an unlocked file.
# 
#     >>> try:
#     ...   raise NotLocked
#     ... except UnlockError:
#     ...   pass
#     
class NotLocked(UnlockError):
    pass

##
# Raised when an attempt is made to unlock a file someone else locked.
# 
#     >>> try:
#     ...   raise NotMyLock
#     ... except UnlockError:
#     ...   pass
#     
class NotMyLock(UnlockError):
    pass

##
# Base class for platform-specific lock classes.
class LockBase:
    ##
    # 
    #         >>> lock = LockBase(_testfile())
    #         
    def __init__(self, path, threaded=True):
        self.path = path
        self.lock_file = os.path.abspath(path) + ".lock"
        self.hostname = socket.gethostname()
        self.pid = os.getpid()
        if threaded:
            tname = "%x-" % thread.get_ident()
        else:
            tname = ""
        dirname = os.path.dirname(self.lock_file)
        self.unique_name = os.path.join(dirname,
                                        "%s.%s%s" % (self.hostname,
                                                     tname,
                                                     self.pid))

    ##
    # 
    #         Acquire the lock.
    # 
    #         * If timeout is omitted (or None), wait forever trying to lock the
    #           file.
    # 
    #         * If timeout > 0, try to acquire the lock for that many seconds.  If
    #           the lock period expires and the file is still locked, raise
    #           LockTimeout.
    # 
    #         * If timeout <= 0, raise AlreadyLocked immediately if the file is
    #           already locked.
    # 
    #         >>> # As simple as it gets.
    #         >>> lock = FileLock(_testfile())
    #         >>> lock.acquire()
    #         >>> lock.release()
    # 
    #         >>> # No timeout test
    #         >>> e1, e2 = threading.Event(), threading.Event()
    #         >>> t = _in_thread(_lock_wait_unlock, e1, e2)
    #         >>> e1.wait()         # wait for thread t to acquire lock
    #         >>> lock2 = FileLock(_testfile())
    #         >>> lock2.is_locked()
    #         True
    #         >>> lock2.i_am_locking()
    #         False
    #         >>> try:
    #         ...   lock2.acquire(timeout=-1)
    #         ... except AlreadyLocked:
    #         ...   pass
    #         ... except Exception, e:
    #         ...   print 'unexpected exception', repr(e)
    #         ... else:
    #         ...   print 'thread', threading.currentThread().getName(),
    #         ...   print 'erroneously locked an already locked file.'
    #         ...   lock2.release()
    #         ...
    #         >>> e2.set()          # tell thread t to release lock
    #         >>> t.join()
    # 
    #         >>> # Timeout test
    #         >>> e1, e2 = threading.Event(), threading.Event()
    #         >>> t = _in_thread(_lock_wait_unlock, e1, e2)
    #         >>> e1.wait() # wait for thread t to acquire filelock
    #         >>> lock2 = FileLock(_testfile())
    #         >>> lock2.is_locked()
    #         True
    #         >>> try:
    #         ...   lock2.acquire(timeout=0.1)
    #         ... except LockTimeout:
    #         ...   pass
    #         ... except Exception, e:
    #         ...   print 'unexpected exception', repr(e)
    #         ... else:
    #         ...   lock2.release()
    #         ...   print 'thread', threading.currentThread().getName(),
    #         ...   print 'erroneously locked an already locked file.'
    #         ...
    #         >>> e2.set()
    #         >>> t.join()
    #         
    def acquire(self, timeout=None):
        pass

    ##
    # 
    #         Release the lock.
    # 
    #         If the file is not locked, raise NotLocked.
    #         >>> lock = FileLock(_testfile())
    #         >>> lock.acquire()
    #         >>> lock.release()
    #         >>> lock.is_locked()
    #         False
    #         >>> lock.i_am_locking()
    #         False
    #         >>> try:
    #         ...   lock.release()
    #         ... except NotLocked:
    #         ...   pass
    #         ... except NotMyLock:
    #         ...   print 'unexpected exception', NotMyLock
    #         ... except Exception, e:
    #         ...   print 'unexpected exception', repr(e)
    #         ... else:
    #         ...   print 'erroneously unlocked file'
    # 
    #         >>> e1, e2 = threading.Event(), threading.Event()
    #         >>> t = _in_thread(_lock_wait_unlock, e1, e2)
    #         >>> e1.wait()
    #         >>> lock2 = FileLock(_testfile())
    #         >>> lock2.is_locked()
    #         True
    #         >>> lock2.i_am_locking()
    #         False
    #         >>> try:
    #         ...   lock2.release()
    #         ... except NotMyLock:
    #         ...   pass
    #         ... except Exception, e:
    #         ...   print 'unexpected exception', repr(e)
    #         ... else:
    #         ...   print 'erroneously unlocked a file locked by another thread.'
    #         ...
    #         >>> e2.set()
    #         >>> t.join()
    #         
    def release(self):
        pass

    ##
    # 
    #         Tell whether or not the file is locked.
    #         >>> lock = FileLock(_testfile())
    #         >>> lock.acquire()
    #         >>> lock.is_locked()
    #         True
    #         >>> lock.release()
    #         >>> lock.is_locked()
    #         False
    #         
    def is_locked(self):
        pass

    ##
    # Return True if this object is locking the file.
    # 
    #         >>> lock1 = FileLock(_testfile(), threaded=False)
    #         >>> lock1.acquire()
    #         >>> lock2 = FileLock(_testfile())
    #         >>> lock1.i_am_locking()
    #         True
    #         >>> lock2.i_am_locking()
    #         False
    #   >>> try:
    #   ...   lock2.acquire(timeout=2)
    #   ... except LockTimeout:
    #         ...   lock2.break_lock()
    #         ...   lock2.is_locked()
    #         ...   lock1.is_locked()
    #         ...   lock2.acquire()
    #         ... else:
    #         ...   print 'expected LockTimeout...'
    #         ...
    #         False
    #         False
    #         >>> lock1.i_am_locking()
    #         False
    #         >>> lock2.i_am_locking()
    #         True
    #         >>> lock2.release()
    #         
    def i_am_locking(self):
        pass

    ##
    # Remove a lock.  Useful if a locking thread failed to unlock.
    # 
    #         >>> lock = FileLock(_testfile())
    #         >>> lock.acquire()
    #         >>> lock2 = FileLock(_testfile())
    #         >>> lock2.is_locked()
    #         True
    #         >>> lock2.break_lock()
    #         >>> lock2.is_locked()
    #         False
    #         >>> try:
    #         ...   lock.release()
    #         ... except NotLocked:
    #         ...   pass
    #         ... except Exception, e:
    #         ...   print 'unexpected exception', repr(e)
    #         ... else:
    #         ...   print 'break lock failed'
    #         
    def break_lock(self):
        pass

    ##
    # Context manager support.
    # 
    #         >>> lock = FileLock(_testfile())
    #         >>> with lock:
    #         ...   lock.is_locked()
    #         ...
    #         True
    #         >>> lock.is_locked()
    #         False
    #         
    def __enter__(self):
        self.acquire()
        return self

    ##
    # Context manager support.
    # 
    #         >>> 'tested in __enter__'
    #         'tested in __enter__'
    #         
    def __exit__(self, *_exc):
        self.release()

##
# Lock access to a file using atomic property of link(2).
class LinkFileLock(LockBase):

    ##
    # 
    #         >>> d = _testfile()
    #         >>> os.mkdir(d)
    #         >>> os.chmod(d, 0444)
    #         >>> try:
    #         ...   lock = LinkFileLock(os.path.join(d, 'test'))
    #         ...   try:
    #         ...     lock.acquire()
    #         ...   except LockFailed:
    #         ...     pass
    #         ...   else:
    #         ...     lock.release()
    #         ...     print 'erroneously locked', os.path.join(d, 'test')
    #         ... finally:
    #         ...   os.chmod(d, 0664)
    #         ...   os.rmdir(d)
    #         
    def acquire(self, timeout=None):
        try:
            open(self.unique_name, "wb").close()
        except IOError:
            raise LockFailed

        end_time = time.time()
        if timeout is not None and timeout > 0:
            end_time += timeout

        while True:
            # Try and create a hard link to it.
            try:
                os.link(self.unique_name, self.lock_file)
            except OSError:
                # Link creation failed.  Maybe we've double-locked?
                nlinks = os.stat(self.unique_name).st_nlink
                if nlinks == 2:
                    # The original link plus the one I created == 2.  We're
                    # good to go.
                    return
                else:
                    # Otherwise the lock creation failed.
                    if timeout is not None and time.time() > end_time:
                        os.unlink(self.unique_name)
                        if timeout > 0:
                            raise LockTimeout
                        else:
                            raise AlreadyLocked
                    time.sleep(timeout is not None and timeout/10 or 0.1)
            else:
                # Link creation succeeded.  We're good to go.
                return

    def release(self):
        if not self.is_locked():
            raise NotLocked
        elif not os.path.exists(self.unique_name):
            raise NotMyLock
        os.unlink(self.unique_name)
        os.unlink(self.lock_file)

    def is_locked(self):
        return os.path.exists(self.lock_file)

    def i_am_locking(self):
        return (self.is_locked() and
                os.path.exists(self.unique_name) and
                os.stat(self.unique_name).st_nlink == 2)

    def break_lock(self):
        if os.path.exists(self.lock_file):
            os.unlink(self.lock_file)

##
# Lock file by creating a directory.
class MkdirFileLock(LockBase):
    ##
    # 
    #         >>> lock = MkdirFileLock(_testfile())
    #         
    def __init__(self, path, threaded=True):
        LockBase.__init__(self, path)
        if threaded:
            tname = "%x-" % thread.get_ident()
        else:
            tname = ""
        # Lock file itself is a directory.  Place the unique file name into
        # it.
        self.unique_name  = os.path.join(self.lock_file,
                                         "%s.%s%s" % (self.hostname,
                                                      tname,
                                                      self.pid))

    def acquire(self, timeout=None):
        end_time = time.time()
        if timeout is not None and timeout > 0:
            end_time += timeout

        if timeout is None:
            wait = 0.1
        else:
            wait = max(0, timeout / 10)

        while True:
            try:
                os.mkdir(self.lock_file)
            except OSError, err:
                if err.errno == errno.EEXIST:
                    # Already locked.
                    if os.path.exists(self.unique_name):
                        # Already locked by me.
                        return
                    if timeout is not None and time.time() > end_time:
                        if timeout > 0:
                            raise LockTimeout
                        else:
                            # Someone else has the lock.
                            raise AlreadyLocked
                    time.sleep(wait)
                else:
                    # Couldn't create the lock for some other reason
                    raise LockFailed
            else:
                open(self.unique_name, "wb").close()
                return

    def release(self):
        if not self.is_locked():
            raise NotLocked
        elif not os.path.exists(self.unique_name):
            raise NotMyLock
        os.unlink(self.unique_name)
        os.rmdir(self.lock_file)

    def is_locked(self):
        return os.path.exists(self.lock_file)

    def i_am_locking(self):
        return (self.is_locked() and
                os.path.exists(self.unique_name))

    def break_lock(self):
        if os.path.exists(self.lock_file):
            for name in os.listdir(self.lock_file):
                os.unlink(os.path.join(self.lock_file, name))
            os.rmdir(self.lock_file)

class SQLiteFileLock(LockBase):
    "Demonstration of using same SQL-based locking."

    import tempfile
    _fd, testdb = tempfile.mkstemp()
    os.close(_fd)
    os.unlink(testdb)
    del _fd, tempfile

    def __init__(self, path, threaded=True):
        LockBase.__init__(self, path, threaded)
        self.lock_file = unicode(self.lock_file)
        self.unique_name = unicode(self.unique_name)

        import sqlite3
        self.connection = sqlite3.connect(SQLiteFileLock.testdb)
        
        c = self.connection.cursor()
        try:
            c.execute("create table locks"
                      "("
                      "   lock_file varchar(32),"
                      "   unique_name varchar(32)"
                      ")")
        except sqlite3.OperationalError:
            pass
        else:
            self.connection.commit()
            import atexit
            atexit.register(os.unlink, SQLiteFileLock.testdb)

    def acquire(self, timeout=None):
        end_time = time.time()
        if timeout is not None and timeout > 0:
            end_time += timeout

        if timeout is None:
            wait = 0.1
        elif timeout <= 0:
            wait = 0
        else:
            wait = timeout / 10

        cursor = self.connection.cursor()

        while True:
            if not self.is_locked():
                # Not locked.  Try to lock it.
                cursor.execute("insert into locks"
                               "  (lock_file, unique_name)"
                               "  values"
                               "  (?, ?)",
                               (self.lock_file, self.unique_name))
                self.connection.commit()

                # Check to see if we are the only lock holder.
                cursor.execute("select * from locks"
                               "  where unique_name = ?",
                               (self.unique_name,))
                rows = cursor.fetchall()
                if len(rows) > 1:
                    # Nope.  Someone else got there.  Remove our lock.
                    cursor.execute("delete from locks"
                                   "  where unique_name = ?",
                                   (self.unique_name,))
                    self.connection.commit()
                else:
                    # Yup.  We're done, so go home.
                    return
            else:
                # Check to see if we are the only lock holder.
                cursor.execute("select * from locks"
                               "  where unique_name = ?",
                               (self.unique_name,))
                rows = cursor.fetchall()
                if len(rows) == 1:
                    # We're the locker, so go home.
                    return
                    
            # Maybe we should wait a bit longer.
            if timeout is not None and time.time() > end_time:
                if timeout > 0:
                    # No more waiting.
                    raise LockTimeout
                else:
                    # Someone else has the lock and we are impatient..
                    raise AlreadyLocked

            # Well, okay.  We'll give it a bit longer.
            time.sleep(wait)

    def release(self):
        if not self.is_locked():
            raise NotLocked
        if not self.i_am_locking():
            raise NotMyLock, ("locker:", self._who_is_locking(),
                              "me:", self.unique_name)
        cursor = self.connection.cursor()
        cursor.execute("delete from locks"
                       "  where unique_name = ?",
                       (self.unique_name,))
        self.connection.commit()

    def _who_is_locking(self):
        cursor = self.connection.cursor()
        cursor.execute("select unique_name from locks"
                       "  where lock_file = ?",
                       (self.lock_file,))
        return cursor.fetchone()[0]
        
    def is_locked(self):
        cursor = self.connection.cursor()
        cursor.execute("select * from locks"
                       "  where lock_file = ?",
                       (self.lock_file,))
        rows = cursor.fetchall()
        return not not rows

    def i_am_locking(self):
        cursor = self.connection.cursor()
        cursor.execute("select * from locks"
                       "  where lock_file = ?"
                       "    and unique_name = ?",
                       (self.lock_file, self.unique_name))
        return not not cursor.fetchall()

    def break_lock(self):
        cursor = self.connection.cursor()
        cursor.execute("delete from locks"
                       "  where lock_file = ?",
                       (self.lock_file,))
        self.connection.commit()

if hasattr(os, "link"):
    FileLock = LinkFileLock
else:
    FileLock = MkdirFileLock

##
# Execute func(*args, **kwargs) after dt seconds.
# 
#     Helper for docttests.
#     
def _in_thread(func, *args, **kwargs):
    def _f():
        func(*args, **kwargs)
    t = threading.Thread(target=_f, name='/*/*')
    t.start()
    return t

##
# Return platform-appropriate lock file name.
# 
#     Helper for doctests.
#     
def _testfile():
    import tempfile
    return os.path.join(tempfile.gettempdir(), 'trash-%s' % os.getpid())

##
# Lock from another thread.
# 
#     Helper for doctests.
#     
def _lock_wait_unlock(event1, event2):
    lock = FileLock(_testfile())
    with lock:
        event1.set()  # we're in,
        event2.wait() # wait for boss's permission to leave

def _test():
    global FileLock

    import doctest
    import sys

    def test_object(c):
        nfailed = ntests = 0
        for (obj, recurse) in ((c, True),
                               (LockBase, True),
                               (sys.modules["__main__"], False)):
            tests = doctest.DocTestFinder(recurse=recurse).find(obj)
            runner = doctest.DocTestRunner(verbose="-v" in sys.argv)
            tests.sort(key = lambda test: test.name)
            for test in tests:
                f, t = runner.run(test)
                nfailed += f
                ntests += t
        print FileLock.__name__, "tests:", ntests, "failed:", nfailed
        return nfailed, ntests

    nfailed = ntests = 0

    if hasattr(os, "link"):
        FileLock = LinkFileLock
        f, t = test_object(FileLock)
        nfailed += f
        ntests += t

    if hasattr(os, "mkdir"):
        FileLock = MkdirFileLock
        f, t = test_object(FileLock)
        nfailed += f
        ntests += t

    try:
        import sqlite3
    except ImportError:
        print "SQLite3 is unavailable - not testing SQLiteFileLock."
    else:
        print "Testing SQLiteFileLock with sqlite", sqlite3.sqlite_version,
        print "& pysqlite", sqlite3.version
        FileLock = SQLiteFileLock
        f, t = test_object(FileLock)
        nfailed += f
        ntests += t

    print "total tests:", ntests, "total failed:", nfailed

if __name__ == "__main__":
    _test()