summaryrefslogtreecommitdiff
path: root/lockfile.py
blob: ee8c2559606b0fc9488d3c2bc5ed772ec70254a0 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
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