1import os
2import sys
3import errno
4import atexit
5import signal
6import logging
7import tempfile
8from .utils import (
9    determine_pid_directory,
10    effective_access,
11)
12try:
13    from contextlib import ContextDecorator as BaseObject
14except ImportError:
15    BaseObject = object
16
17
18DEFAULT_PID_DIR = determine_pid_directory()
19DEFAULT_CHMOD = 0o644
20PID_CHECK_EMPTY = "PID_CHECK_EMPTY"
21PID_CHECK_NOFILE = "PID_CHECK_NOFILE"
22PID_CHECK_SAMEPID = "PID_CHECK_SAMEPID"
23PID_CHECK_NOTRUNNING = "PID_CHECK_NOTRUNNING"
24
25
26class PidFileError(Exception):
27    pass
28
29
30class PidFileConfigurationError(Exception):
31    pass
32
33
34class PidFileUnreadableError(PidFileError):
35    pass
36
37
38class PidFileAlreadyRunningError(PidFileError):
39    def __init__(self, message, pid=None):
40        self.message = message
41        self.pid = pid
42
43
44class PidFileAlreadyLockedError(PidFileError):
45    pass
46
47
48class PidFileBase(BaseObject):
49    __slots__ = (
50        "pid", "pidname", "piddir", "enforce_dotpid_postfix",
51        "register_term_signal_handler", "register_atexit", "filename",
52        "fh", "lock_pidfile", "chmod", "uid", "gid", "force_tmpdir",
53        "allow_samepid", "_logger", "_is_setup", "_need_cleanup",
54    )
55
56    def __init__(self, pidname=None, piddir=None, enforce_dotpid_postfix=True,
57                 register_term_signal_handler="auto", register_atexit=True,
58                 lock_pidfile=True, chmod=DEFAULT_CHMOD, uid=-1, gid=-1, force_tmpdir=False,
59                 allow_samepid=False):
60        self.pidname = pidname
61        self.piddir = piddir
62        self.enforce_dotpid_postfix = enforce_dotpid_postfix
63        self.register_term_signal_handler = register_term_signal_handler
64        self.register_atexit = register_atexit
65        self.lock_pidfile = lock_pidfile
66        self.chmod = chmod
67        self.uid = uid
68        self.gid = gid
69        self.force_tmpdir = force_tmpdir
70        self.allow_samepid = allow_samepid
71
72        self.fh = None
73        self.filename = None
74        self.pid = None
75
76        self._logger = None
77        self._is_setup = False
78        self._need_cleanup = False
79
80    @property
81    def logger(self):
82        if not self._logger:
83            self._logger = logging.getLogger("PidFile")
84
85        return self._logger
86
87    def setup(self):
88        if not self._is_setup:
89            self.logger.debug("%r entering setup", self)
90            if self.filename is None:
91                self.pid = os.getpid()
92                self.filename = self._make_filename()
93                self._register_term_signal()
94
95            if self.register_atexit:
96                atexit.register(self.close)
97
98            # setup should only be performed once
99            self._is_setup = True
100
101    def _make_filename(self):
102        pidname = self.pidname
103        piddir = self.piddir
104        if pidname is None:
105            pidname = "%s.pid" % os.path.basename(sys.argv[0])
106        if self.enforce_dotpid_postfix and not pidname.endswith(".pid"):
107            pidname = "%s.pid" % pidname
108        if piddir is None:
109            if os.path.isdir(DEFAULT_PID_DIR) and self.force_tmpdir is False:
110                piddir = DEFAULT_PID_DIR
111            else:
112                piddir = tempfile.gettempdir()
113
114        if os.path.exists(piddir) and not os.path.isdir(piddir):
115            raise IOError("Pid file directory '%s' exists but is not a directory" % piddir)
116        if not os.path.isdir(piddir):
117            os.makedirs(piddir)
118        if not effective_access(piddir, os.R_OK):
119            raise IOError("Pid file directory '%s' cannot be read" % piddir)
120        if not effective_access(piddir, os.W_OK | os.X_OK):
121            raise IOError("Pid file directory '%s' cannot be written to" % piddir)
122
123        filename = os.path.abspath(os.path.join(piddir, pidname))
124        return filename
125
126    def _register_term_signal(self):
127        register_term_signal_handler = self.register_term_signal_handler
128        if register_term_signal_handler == "auto":
129            if signal.getsignal(signal.SIGTERM) == signal.SIG_DFL:
130                register_term_signal_handler = True
131            else:
132                register_term_signal_handler = False
133
134        if callable(register_term_signal_handler):
135            signal.signal(signal.SIGTERM, register_term_signal_handler)
136        elif register_term_signal_handler:
137            # Register TERM signal handler to make sure atexit runs on TERM signal
138            def sigterm_noop_handler(*args, **kwargs):
139                raise SystemExit(1)
140
141            signal.signal(signal.SIGTERM, sigterm_noop_handler)
142
143    def _inner_check(self, fh):
144        try:
145            fh.seek(0)
146            pid_str = fh.read(16).split("\n", 1)[0].strip()
147            if not pid_str:
148                return PID_CHECK_EMPTY
149            pid = int(pid_str)
150        except (IOError, ValueError) as exc:
151            self.close(fh=fh)
152            raise PidFileUnreadableError(exc)
153        else:
154            if self.allow_samepid and self.pid == pid:
155                return PID_CHECK_SAMEPID
156
157        try:
158            if self._pid_exists(pid):
159                raise PidFileAlreadyRunningError("Program already running with pid: %d" % pid, pid=pid)
160            else:
161                return PID_CHECK_NOTRUNNING
162        except PidFileAlreadyRunningError:
163            self.close(fh=fh, cleanup=False)
164            raise
165
166    def _pid_exists(self, pid):
167        raise NotImplementedError()
168
169    def _flock(self, fileno):
170        raise NotImplementedError()
171
172    def _chmod(self):
173        raise NotImplementedError()
174
175    def _chown(self):
176        raise NotImplementedError()
177
178    def check(self):
179        self.setup()
180
181        self.logger.debug("%r check pidfile: %s", self, self.filename)
182
183        if self.fh is None:
184            if self.filename and os.path.isfile(self.filename):
185                with open(self.filename, "r") as fh:
186                    return self._inner_check(fh)
187            return PID_CHECK_NOFILE
188
189        return self._inner_check(self.fh)
190
191    def create(self):
192        self.setup()
193
194        self.logger.debug("%r create pidfile: %s", self, self.filename)
195        self.fh = open(self.filename, "a+")
196        if self.lock_pidfile:
197            try:
198                self._flock(self.fh.fileno())
199            except IOError as exc:
200                if not self.allow_samepid:
201                    self.close(cleanup=False)
202                    raise PidFileAlreadyLockedError(exc)
203
204        check_result = self.check()
205        if check_result == PID_CHECK_SAMEPID:
206            return
207
208        self._chmod()
209        self._chown()
210
211        self.fh.seek(0)
212        self.fh.truncate()
213        # pidfile must be composed of the pid number and a newline character
214        self.fh.write("%d\n" % self.pid)
215        self.fh.flush()
216        self.fh.seek(0)
217        self._need_cleanup = True
218
219    def close(self, fh=None, cleanup=None):
220        self.logger.debug("%r closing pidfile: %s", self, self.filename)
221        cleanup = self._need_cleanup if cleanup is None else cleanup
222
223        if not fh:
224            fh = self.fh
225        try:
226            if fh is None:
227                return
228            fh.close()
229        except IOError as exc:
230            # ignore error when file was already closed
231            if exc.errno != errno.EBADF:
232                raise
233        finally:
234            if self.filename and os.path.isfile(self.filename) and cleanup:
235                os.remove(self.filename)
236                self._need_cleanup = False
237
238    def __enter__(self):
239        self.create()
240        return self
241
242    def __exit__(self, exc_type=None, exc_value=None, exc_tb=None):
243        self.close()
244