| 1 | import psutil |
|---|
| 2 | |
|---|
| 3 | # the docs are a little misleading, but this is either WindowsFileLock |
|---|
| 4 | # or UnixFileLock depending upon the platform we're currently on |
|---|
| 5 | from filelock import FileLock, Timeout |
|---|
| 6 | |
|---|
| 7 | |
|---|
| 8 | class ProcessInTheWay(Exception): |
|---|
| 9 | """ |
|---|
| 10 | our pidfile points at a running process |
|---|
| 11 | """ |
|---|
| 12 | |
|---|
| 13 | |
|---|
| 14 | class InvalidPidFile(Exception): |
|---|
| 15 | """ |
|---|
| 16 | our pidfile isn't well-formed |
|---|
| 17 | """ |
|---|
| 18 | |
|---|
| 19 | |
|---|
| 20 | class CannotRemovePidFile(Exception): |
|---|
| 21 | """ |
|---|
| 22 | something went wrong removing the pidfile |
|---|
| 23 | """ |
|---|
| 24 | |
|---|
| 25 | |
|---|
| 26 | def _pidfile_to_lockpath(pidfile): |
|---|
| 27 | """ |
|---|
| 28 | internal helper. |
|---|
| 29 | :returns FilePath: a path to use for file-locking the given pidfile |
|---|
| 30 | """ |
|---|
| 31 | return pidfile.sibling("{}.lock".format(pidfile.basename())) |
|---|
| 32 | |
|---|
| 33 | |
|---|
| 34 | def parse_pidfile(pidfile): |
|---|
| 35 | """ |
|---|
| 36 | :param FilePath pidfile: |
|---|
| 37 | :returns tuple: 2-tuple of pid, creation-time as int, float |
|---|
| 38 | :raises InvalidPidFile: on error |
|---|
| 39 | """ |
|---|
| 40 | with pidfile.open("r") as f: |
|---|
| 41 | content = f.read().decode("utf8").strip() |
|---|
| 42 | try: |
|---|
| 43 | pid, starttime = content.split() |
|---|
| 44 | pid = int(pid) |
|---|
| 45 | starttime = float(starttime) |
|---|
| 46 | except ValueError: |
|---|
| 47 | raise InvalidPidFile( |
|---|
| 48 | "found invalid PID file in {}".format( |
|---|
| 49 | pidfile |
|---|
| 50 | ) |
|---|
| 51 | ) |
|---|
| 52 | return pid, starttime |
|---|
| 53 | |
|---|
| 54 | |
|---|
| 55 | def check_pid_process(pidfile): |
|---|
| 56 | """ |
|---|
| 57 | If another instance appears to be running already, raise an |
|---|
| 58 | exception. Otherwise, write our PID + start time to the pidfile |
|---|
| 59 | and arrange to delete it upon exit. |
|---|
| 60 | |
|---|
| 61 | :param FilePath pidfile: the file to read/write our PID from. |
|---|
| 62 | |
|---|
| 63 | :raises ProcessInTheWay: if a running process exists at our PID |
|---|
| 64 | """ |
|---|
| 65 | lock_path = _pidfile_to_lockpath(pidfile) |
|---|
| 66 | |
|---|
| 67 | try: |
|---|
| 68 | # a short timeout is fine, this lock should only be active |
|---|
| 69 | # while someone is reading or deleting the pidfile .. and |
|---|
| 70 | # facilitates testing the locking itself. |
|---|
| 71 | with FileLock(lock_path.path, timeout=2): |
|---|
| 72 | # check if we have another instance running already |
|---|
| 73 | if pidfile.exists(): |
|---|
| 74 | pid, starttime = parse_pidfile(pidfile) |
|---|
| 75 | try: |
|---|
| 76 | # if any other process is running at that PID, let the |
|---|
| 77 | # user decide if this is another legitimate |
|---|
| 78 | # instance. Automated programs may use the start-time to |
|---|
| 79 | # help decide this (if the PID is merely recycled, the |
|---|
| 80 | # start-time won't match). |
|---|
| 81 | psutil.Process(pid) |
|---|
| 82 | raise ProcessInTheWay( |
|---|
| 83 | "A process is already running as PID {}".format(pid) |
|---|
| 84 | ) |
|---|
| 85 | except psutil.NoSuchProcess: |
|---|
| 86 | print( |
|---|
| 87 | "'{pidpath}' refers to {pid} that isn't running".format( |
|---|
| 88 | pidpath=pidfile.path, |
|---|
| 89 | pid=pid, |
|---|
| 90 | ) |
|---|
| 91 | ) |
|---|
| 92 | # nothing is running at that PID so it must be a stale file |
|---|
| 93 | pidfile.remove() |
|---|
| 94 | |
|---|
| 95 | # write our PID + start-time to the pid-file |
|---|
| 96 | proc = psutil.Process() |
|---|
| 97 | with pidfile.open("w") as f: |
|---|
| 98 | f.write("{} {}\n".format(proc.pid, proc.create_time()).encode("utf8")) |
|---|
| 99 | except Timeout: |
|---|
| 100 | raise ProcessInTheWay( |
|---|
| 101 | "Another process is still locking {}".format(pidfile.path) |
|---|
| 102 | ) |
|---|
| 103 | |
|---|
| 104 | |
|---|
| 105 | def cleanup_pidfile(pidfile): |
|---|
| 106 | """ |
|---|
| 107 | Remove the pidfile specified (respecting locks). If anything at |
|---|
| 108 | all goes wrong, `CannotRemovePidFile` is raised. |
|---|
| 109 | """ |
|---|
| 110 | lock_path = _pidfile_to_lockpath(pidfile) |
|---|
| 111 | with FileLock(lock_path.path): |
|---|
| 112 | try: |
|---|
| 113 | pidfile.remove() |
|---|
| 114 | except Exception as e: |
|---|
| 115 | raise CannotRemovePidFile( |
|---|
| 116 | "Couldn't remove '{pidfile}': {err}.".format( |
|---|
| 117 | pidfile=pidfile.path, |
|---|
| 118 | err=e, |
|---|
| 119 | ) |
|---|
| 120 | ) |
|---|