Advertisement

← refacto-feature-scout

scout/locks.py

2.3 KB Download file
"""One job at a time on this computer.

Feature Scout can be started by hand while it is already running on its own. Rather than have
two of them building the same thing, the second one notices and waits its turn — or steps aside.
Windows has no `flock`, so the same idea is done with the tool that platform provides.
"""
from __future__ import annotations

import os
import time
from pathlib import Path

from .config import LOCK_PATH

BUSY = ("Feature Scout is already busy with something on this computer. "
        "It'll get to this next.")

if os.name == "nt":                                          # pragma: no cover - Windows only
    import msvcrt

    def _take(handle) -> None:
        msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1)

    def _release(handle) -> None:
        try:
            handle.seek(0)
            msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1)
        except OSError:
            pass
else:
    import fcntl

    def _take(handle) -> None:
        fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)

    def _release(handle) -> None:
        try:
            fcntl.flock(handle, fcntl.LOCK_UN)
        except OSError:
            pass


class Busy(Exception):
    """Something else is already running."""


class Lock:
    """Held for as long as a job runs. `wait_s` is how long to wait for a turn."""

    def __init__(self, wait_s: int = 0, path: Path | None = None,
                 sleep=time.sleep, clock=time.time):
        self.wait_s = wait_s
        self.path = Path(path or LOCK_PATH)
        self.handle = None
        self._sleep = sleep
        self._clock = clock

    def __enter__(self) -> "Lock":
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.handle = open(self.path, "w")
        deadline = self._clock() + self.wait_s
        while True:
            try:
                _take(self.handle)
                return self
            except OSError:
                if self._clock() >= deadline:
                    self.handle.close()
                    self.handle = None
                    raise Busy(BUSY) from None
                self._sleep(15)

    def __exit__(self, *exc) -> None:
        if self.handle:
            _release(self.handle)
            self.handle.close()
            self.handle = None