lockfile.py

   1 #!/usr/bin/env python3
   2
   3 from dataclasses import dataclass
   4 import datetime
   5 import json
   6 import logging
   7 import os
   8 import signal
   9 import sys
  10 from typing import Optional
  11
  12 import decorator_utils
  13
  14
  15 logger = logging.getLogger(__name__)
  16
  17
  18 class LockFileException(Exception):
  19     pass
  20
  21
  22 @dataclass
  23 class LockFileContents:
  24     pid: int
  25     commandline: str
  26     expiration_timestamp: float
  27
  28
  29 class LockFile(object):
  30     """A file locking mechanism that has context-manager support so you
  31     can use it in a with statement.
  32     """
  33
  34     def __init__(
  35             self,
  36             lockfile_path: str,
  37             *,
  38             do_signal_cleanup: bool = True,
  39             expiration_timestamp: Optional[float] = None,
  40     ) -> None:
  41         self.is_locked = False
  42         self.lockfile = lockfile_path
  43         if do_signal_cleanup:
  44             signal.signal(signal.SIGINT, self._signal)
  45             signal.signal(signal.SIGTERM, self._signal)
  46         self.expiration_timestamp = expiration_timestamp
  47
  48     def locked(self):
  49         return self.is_locked
  50
  51     def available(self):
  52         return not os.path.exists(self.lockfile)
  53
  54     def try_acquire_lock_once(self) -> bool:
  55         logger.debug(f"Trying to acquire {self.lockfile}.")
  56         try:
  57             # Attempt to create the lockfile.  These flags cause
  58             # os.open to raise an OSError if the file already
  59             # exists.
  60             fd = os.open(self.lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR)
  61             with os.fdopen(fd, "a") as f:
  62                 contents = self._get_lockfile_contents()
  63                 logger.debug(contents)
  64                 f.write(contents)
  65             logger.debug(f'Success; I own {self.lockfile}.')
  66             self.is_locked = True
  67             return True
  68         except OSError:
  69             pass
  70         logger.debug(f'Failed; I could not acquire {self.lockfile}.')
  71         return False
  72
  73     def acquire_with_retries(
  74             self,
  75             *,
  76             initial_delay: float = 1.0,
  77             backoff_factor: float = 2.0,
  78             max_attempts = 5
  79     ) -> bool:
  80
  81         @decorator_utils.retry_if_false(tries = max_attempts,
  82                                         delay_sec = initial_delay,
  83                                         backoff = backoff_factor)
  84         def _try_acquire_lock_with_retries() -> bool:
  85             success = self.try_acquire_lock_once()
  86             if not success and os.path.exists(self.lockfile):
  87                 self._detect_stale_lockfile()
  88             return success
  89
  90         if os.path.exists(self.lockfile):
  91             self._detect_stale_lockfile()
  92         return _try_acquire_lock_with_retries()
  93
  94     def release(self):
  95         try:
  96             os.unlink(self.lockfile)
  97         except Exception as e:
  98             logger.exception(e)
  99         self.is_locked = False
 100
 101     def __enter__(self):
 102         if self.acquire_with_retries():
 103             return self
 104         msg = f"Couldn't acquire {self.lockfile}; giving up."
 105         logger.warning(msg)
 106         raise LockFileException(msg)
 107
 108     def __exit__(self, type, value, traceback):
 109         self.release()
 110
 111     def __del__(self):
 112         if self.is_locked:
 113             self.release()
 114
 115     def _signal(self, *args):
 116         if self.is_locked:
 117             self.release()
 118
 119     def _get_lockfile_contents(self) -> str:
 120         contents = LockFileContents(
 121             pid = os.getpid(),
 122             commandline = ' '.join(sys.argv),
 123             expiration_timestamp = self.expiration_timestamp
 124         )
 125         return json.dumps(contents.__dict__)
 126
 127     def _detect_stale_lockfile(self) -> None:
 128         try:
 129             with open(self.lockfile, 'r') as rf:
 130                 lines = rf.readlines()
 131                 if len(lines) == 1:
 132                     line = lines[0]
 133                     line_dict = json.loads(line)
 134                     contents = LockFileContents(**line_dict)
 135                     logger.debug(f'Blocking lock contents="{contents}"')
 136
 137                     # Does the PID exist still?
 138                     try:
 139                         os.kill(contents.pid, 0)
 140                     except OSError:
 141                         logger.debug('The pid seems stale; killing the lock.')
 142                         self.release()
 143
 144                     # Has the lock expiration expired?
 145                     if contents.expiration_timestamp is not None:
 146                         now = datetime.datetime.now().timestamp()
 147                         if now > contents.expiration_datetime:
 148                             logger.debug('The expiration time has passed; ' +
 149                                          'killing the lock')
 150                             self.release()
 151         except Exception:
 152             pass