proxy_runtime.py

"""Thread-owned asyncio runtime for the local HTTP proxy."""
import asyncio
import logging
import threading

from proxy import run_proxy


logger = logging.getLogger("magic-proxy.runtime")


class ProxyRuntime:
    """Own one generation of loop/task/thread and stop it before replacement."""

    def __init__(self, stats):
        self._stats = stats
        self._lock = threading.Lock()
        self._generation = 0
        self._thread = None
        self._loop = None
        self._task = None
        self._stop_event = None
        self._error = ""

    @property
    def running(self):
        with self._lock:
            return bool(self._thread and self._thread.is_alive() and not self._stop_event.is_set())

    @property
    def error(self):
        with self._lock:
            return self._error

    def start(self, config):
        self.stop()
        with self._lock:
            self._generation += 1
            generation = self._generation
            stop_event = threading.Event()
            self._stop_event = stop_event
            self._error = ""

        def worker():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            task = None
            try:
                if stop_event.is_set():
                    return
                control = {}
                task = loop.create_task(run_proxy(dict(config), self._stats, control))
                with self._lock:
                    if generation != self._generation:
                        task.cancel()
                    else:
                        self._loop = loop
                        self._task = task
                if stop_event.is_set():
                    task.cancel()
                loop.run_until_complete(task)
            except asyncio.CancelledError:
                pass
            except Exception as exc:
                logger.exception("HTTP proxy generation %d stopped", generation)
                with self._lock:
                    if generation == self._generation:
                        self._error = str(exc)
            finally:
                pending = asyncio.all_tasks(loop)
                for pending_task in pending:
                    pending_task.cancel()
                if pending:
                    loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
                loop.close()
                with self._lock:
                    if generation == self._generation:
                        self._loop = None
                        self._task = None
                        self._thread = None

        thread = threading.Thread(
            target=worker, name=f"MagicProxyHTTP-{generation}", daemon=True,
        )
        with self._lock:
            self._thread = thread
        thread.start()
        return True

    def stop(self, timeout=5):
        with self._lock:
            thread = self._thread
            loop = self._loop
            task = self._task
            stop_event = self._stop_event
        if stop_event:
            stop_event.set()
        if loop and task and loop.is_running():
            try:
                loop.call_soon_threadsafe(task.cancel)
            except RuntimeError:
                pass
        if thread and thread is not threading.current_thread():
            thread.join(timeout=timeout)
        alive = bool(thread and thread.is_alive())
        if alive:
            logger.error("HTTP proxy thread did not stop within %.1fs", timeout)
        return not alive