diff options
Diffstat (limited to 'lib/spack/llnl/util/lock.py')
-rw-r--r-- | lib/spack/llnl/util/lock.py | 89 |
1 files changed, 79 insertions, 10 deletions
diff --git a/lib/spack/llnl/util/lock.py b/lib/spack/llnl/util/lock.py index 479a1b0167..4a4aec2385 100644 --- a/lib/spack/llnl/util/lock.py +++ b/lib/spack/llnl/util/lock.py @@ -28,6 +28,9 @@ import errno import time import socket +__all__ = ['Lock', 'LockTransaction', 'WriteTransaction', 'ReadTransaction', + 'LockError'] + # Default timeout in seconds, after which locks will raise exceptions. _default_timeout = 60 @@ -36,13 +39,20 @@ _sleep_time = 1e-5 class Lock(object): - def __init__(self,file_path): + """This is an implementation of a filesystem lock using Python's lockf. + + In Python, `lockf` actually calls `fcntl`, so this should work with any + filesystem implementation that supports locking through the fcntl calls. + This includes distributed filesystems like Lustre (when flock is enabled) + and recent NFS versions. + + """ + def __init__(self, file_path): self._file_path = file_path self._fd = None self._reads = 0 self._writes = 0 - def _lock(self, op, timeout): """This takes a lock using POSIX locks (``fnctl.lockf``). @@ -63,7 +73,9 @@ class Lock(object): fcntl.lockf(self._fd, op | fcntl.LOCK_NB) if op == fcntl.LOCK_EX: - os.write(self._fd, "pid=%s,host=%s" % (os.getpid(), socket.getfqdn())) + os.write( + self._fd, + "pid=%s,host=%s" % (os.getpid(), socket.getfqdn())) return except IOError as error: @@ -75,7 +87,6 @@ class Lock(object): raise LockError("Timed out waiting for lock.") - def _unlock(self): """Releases a lock using POSIX locks (``fcntl.lockf``) @@ -83,11 +94,10 @@ class Lock(object): be masquerading as write locks, but this removes either. """ - fcntl.lockf(self._fd,fcntl.LOCK_UN) + fcntl.lockf(self._fd, fcntl.LOCK_UN) os.close(self._fd) self._fd = None - def acquire_read(self, timeout=_default_timeout): """Acquires a recursive, shared lock for reading. @@ -107,7 +117,6 @@ class Lock(object): self._reads += 1 return False - def acquire_write(self, timeout=_default_timeout): """Acquires a recursive, exclusive lock for writing. @@ -127,7 +136,6 @@ class Lock(object): self._writes += 1 return False - def release_read(self): """Releases a read lock. @@ -148,7 +156,6 @@ class Lock(object): self._reads -= 1 return False - def release_write(self): """Releases a write lock. @@ -170,6 +177,68 @@ class Lock(object): return False +class LockTransaction(object): + """Simple nested transaction context manager that uses a file lock. + + This class can trigger actions when the lock is acquired for the + first time and released for the last. + + If the acquire_fn returns a value, it is used as the return value for + __enter__, allowing it to be passed as the `as` argument of a `with` + statement. + + If acquire_fn returns a context manager, *its* `__enter__` function will be + called in `__enter__` after acquire_fn, and its `__exit__` funciton will be + called before `release_fn` in `__exit__`, allowing you to nest a context + manager to be used along with the lock. + + Timeout for lock is customizable. + + """ + + def __init__(self, lock, acquire_fn=None, release_fn=None, + timeout=_default_timeout): + self._lock = lock + self._timeout = timeout + self._acquire_fn = acquire_fn + self._release_fn = release_fn + self._as = None + + def __enter__(self): + if self._enter() and self._acquire_fn: + self._as = self._acquire_fn() + if hasattr(self._as, '__enter__'): + return self._as.__enter__() + else: + return self._as + + def __exit__(self, type, value, traceback): + suppress = False + if self._exit(): + if self._as and hasattr(self._as, '__exit__'): + if self._as.__exit__(type, value, traceback): + suppress = True + if self._release_fn: + if self._release_fn(type, value, traceback): + suppress = True + return suppress + + +class ReadTransaction(LockTransaction): + def _enter(self): + return self._lock.acquire_read(self._timeout) + + def _exit(self): + return self._lock.release_read() + + +class WriteTransaction(LockTransaction): + def _enter(self): + return self._lock.acquire_write(self._timeout) + + def _exit(self): + return self._lock.release_write() + + class LockError(Exception): """Raised when an attempt to acquire a lock times out.""" - pass |