devplacepy/devplacepy/services/background.py

# retoor <retoor@molodetz.nl>

import asyncio
import logging
import queue
import threading
from typing import Any, Callable, Optional

logger = logging.getLogger(__name__)

QUEUE_MAXSIZE = 10000
JOIN_TIMEOUT_SECONDS = 10
GET_TIMEOUT_SECONDS = 1


class _Stop:
    pass


_STOP = _Stop()


class BackgroundQueue:
    def __init__(self, maxsize: int = QUEUE_MAXSIZE) -> None:
        self._queue: "queue.Queue" = queue.Queue(maxsize=maxsize)
        self._thread: Optional[threading.Thread] = None
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._counters_lock = threading.Lock()
        self._submitted = 0
        self._processed = 0
        self._failed = 0
        self._inline = 0

    @property
    def running(self) -> bool:
        return self._thread is not None and self._thread.is_alive()

    @property
    def loop(self) -> Optional[asyncio.AbstractEventLoop]:
        return self._loop

    def submit(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> None:
        with self._counters_lock:
            self._submitted += 1
        if not self.running:
            self._execute(fn, args, kwargs, inline=True)
            return
        try:
            self._queue.put_nowait((fn, args, kwargs))
        except queue.Full:
            logger.warning("background queue full; running task inline")
            self._execute(fn, args, kwargs, inline=True)

    def _execute(self, fn: Callable[..., Any], args: tuple, kwargs: dict, inline: bool = False) -> None:
        try:
            fn(*args, **kwargs)
            with self._counters_lock:
                self._processed += 1
                if inline:
                    self._inline += 1
        except Exception as exc:
            with self._counters_lock:
                self._failed += 1
                if inline:
                    self._inline += 1
            logger.warning("background task %s failed: %s", getattr(fn, "__name__", fn), exc)
        finally:
            self._release_db_lock()

    def _release_db_lock(self) -> None:
        try:
            from devplacepy.database import refresh_snapshot

            refresh_snapshot()
        except Exception as exc:
            logger.warning("background queue could not release db lock: %s", exc)

    async def start(self) -> None:
        if self.running:
            return
        self._loop = asyncio.get_running_loop()
        self._thread = threading.Thread(
            target=self._worker, name="background-queue", daemon=True
        )
        self._thread.start()
        logger.info("background task queue worker thread started")

    async def stop(self) -> None:
        thread = self._thread
        self._thread = None
        self._loop = None
        if thread is not None:
            self._queue.put(_STOP)
            await asyncio.to_thread(thread.join, JOIN_TIMEOUT_SECONDS)
        self._flush_remaining()
        logger.info(
            "background task queue stopped processed=%s failed=%s inline=%s",
            self._processed,
            self._failed,
            self._inline,
        )

    def _flush_remaining(self) -> None:
        while True:
            try:
                item = self._queue.get_nowait()
            except queue.Empty:
                break
            if item is _STOP:
                continue
            fn, args, kwargs = item
            self._execute(fn, args, kwargs)

    def _worker(self) -> None:
        while True:
            try:
                item = self._queue.get(timeout=GET_TIMEOUT_SECONDS)
            except queue.Empty:
                continue
            if item is _STOP:
                return
            fn, args, kwargs = item
            self._execute(fn, args, kwargs)

    def stats(self) -> dict:
        return {
            "running": self.running,
            "depth": self._queue.qsize(),
            "submitted": self._submitted,
            "processed": self._processed,
            "failed": self._failed,
            "inline": self._inline,
        }


background = BackgroundQueue()