1 files changed, 152 insertions, 0 deletions
diff --git a/lockfile.py b/lockfile.py
new file mode 100644
index 0000000..ee8c255
--- /dev/null
+++ b/lockfile.py
@@ -0,0 +1,152 @@
+#!/usr/bin/env python3
+
+from dataclasses import dataclass
+import datetime
+import json
+import logging
+import os
+import signal
+import sys
+from typing import Optional
+
+import decorator_utils
+
+
+logger = logging.getLogger(__name__)
+
+
+class LockFileException(Exception):
+    pass
+
+
+@dataclass
+class LockFileContents:
+    pid: int
+    commandline: str
+    expiration_timestamp: float
+
+
+class LockFile(object):
+    """A file locking mechanism that has context-manager support so you
+    can use it in a with statement.
+    """
+
+    def __init__(
+            self,
+            lockfile_path: str,
+            *,
+            do_signal_cleanup: bool = True,
+            expiration_timestamp: Optional[float] = None,
+    ) -> None:
+        self.is_locked = False
+        self.lockfile = lockfile_path
+        if do_signal_cleanup:
+            signal.signal(signal.SIGINT, self._signal)
+            signal.signal(signal.SIGTERM, self._signal)
+        self.expiration_timestamp = expiration_timestamp
+
+    def locked(self):
+        return self.is_locked
+
+    def available(self):
+        return not os.path.exists(self.lockfile)
+
+    def try_acquire_lock_once(self) -> bool:
+        logger.debug(f"Trying to acquire {self.lockfile}.")
+        try:
+            # Attempt to create the lockfile.  These flags cause
+            # os.open to raise an OSError if the file already
+            # exists.
+            fd = os.open(self.lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR)
+            with os.fdopen(fd, "a") as f:
+                contents = self._get_lockfile_contents()
+                logger.debug(contents)
+                f.write(contents)
+            logger.debug(f'Success; I own {self.lockfile}.')
+            self.is_locked = True
+            return True
+        except OSError:
+            pass
+        logger.debug(f'Failed; I could not acquire {self.lockfile}.')
+        return False
+
+    def acquire_with_retries(
+            self,
+            *,
+            initial_delay: float = 1.0,
+            backoff_factor: float = 2.0,
+            max_attempts = 5
+    ) -> bool:
+
+        @decorator_utils.retry_if_false(tries = max_attempts,
+                                        delay_sec = initial_delay,
+                                        backoff = backoff_factor)
+        def _try_acquire_lock_with_retries() -> bool:
+            success = self.try_acquire_lock_once()
+            if not success and os.path.exists(self.lockfile):
+                self._detect_stale_lockfile()
+            return success
+
+        if os.path.exists(self.lockfile):
+            self._detect_stale_lockfile()
+        return _try_acquire_lock_with_retries()
+
+    def release(self):
+        try:
+            os.unlink(self.lockfile)
+        except Exception as e:
+            logger.exception(e)
+        self.is_locked = False
+
+    def __enter__(self):
+        if self.acquire_with_retries():
+            return self
+        msg = f"Couldn't acquire {self.lockfile}; giving up."
+        logger.warning(msg)
+        raise LockFileException(msg)
+
+    def __exit__(self, type, value, traceback):
+        self.release()
+
+    def __del__(self):
+        if self.is_locked:
+            self.release()
+
+    def _signal(self, *args):
+        if self.is_locked:
+            self.release()
+
+    def _get_lockfile_contents(self) -> str:
+        contents = LockFileContents(
+            pid = os.getpid(),
+            commandline = ' '.join(sys.argv),
+            expiration_timestamp = self.expiration_timestamp
+        )
+        return json.dumps(contents.__dict__)
+
+    def _detect_stale_lockfile(self) -> None:
+        try:
+            with open(self.lockfile, 'r') as rf:
+                lines = rf.readlines()
+                if len(lines) == 1:
+                    line = lines[0]
+                    line_dict = json.loads(line)
+                    contents = LockFileContents(**line_dict)
+                    logger.debug(f'Blocking lock contents="{contents}"')
+
+                    # Does the PID exist still?
+                    try:
+                        os.kill(contents.pid, 0)
+                    except OSError:
+                        logger.debug('The pid seems stale; killing the lock.')
+                        self.release()
+
+                    # Has the lock expiration expired?
+                    if contents.expiration_timestamp is not None:
+                        now = datetime.datetime.now().timestamp()
+                        if now > contents.expiration_datetime:
+                            logger.debug('The expiration time has passed; ' +
+                                         'killing the lock')
+                            self.release()
+        except Exception:
+            pass