queue.py

"""Manage the communication sessions towards the USB-Stick."""

from __future__ import annotations

from asyncio import PriorityQueue, Task, get_running_loop, sleep
from collections.abc import Callable
from dataclasses import dataclass
import logging

from ..api import StickEvent
from ..constants import REPORT_QUEUE_FILLING_UP
from ..exceptions import MessageError, NodeTimeout, StickError, StickTimeout
from ..messages import Priority
from ..messages.requests import NodePingRequest, PlugwiseCancelRequest, PlugwiseRequest
from ..messages.responses import PlugwiseResponse
from .manager import StickConnectionManager

_LOGGER = logging.getLogger(__name__)


@dataclass
class RequestState:
    """Node hardware information."""

    session: bytes
    zigbee_address: int


class StickQueue:
    """Manage queue of all request sessions."""

    def __init__(self) -> None:
        """Initialize the message session controller."""
        self._stick: StickConnectionManager | None = None
        self._loop = get_running_loop()
        self._submit_queue: PriorityQueue[PlugwiseRequest] = PriorityQueue()
        self._submit_worker_task: Task[None] | None = None
        self._unsubscribe_connection_events: Callable[[], None] | None = None
        self._running = False

    @property
    def is_running(self) -> bool:
        """Return the state of the queue."""
        return self._running

    def start(self, stick_connection_manager: StickConnectionManager) -> None:
        """Start sending request from queue."""
        if self._running:
            raise StickError("Cannot start queue manager, already running")
        self._stick = stick_connection_manager
        if self._stick.is_connected:
            self._running = True
        self._unsubscribe_connection_events = self._stick.subscribe_to_stick_events(
            self._handle_stick_event, (StickEvent.CONNECTED, StickEvent.DISCONNECTED)
        )

    async def _handle_stick_event(self, event: StickEvent) -> None:
        """Handle events from stick."""
        if event is StickEvent.CONNECTED:
            self._running = True
        elif event is StickEvent.DISCONNECTED:
            self._running = False

    async def stop(self) -> None:
        """Stop sending from queue."""
        _LOGGER.debug("Stop queue")
        if self._unsubscribe_connection_events is not None:
            self._unsubscribe_connection_events()
        self._running = False
        if self._submit_worker_task is not None and not self._submit_worker_task.done():
            cancel_request = PlugwiseCancelRequest()
            await self._submit_queue.put(cancel_request)
            await self._submit_worker_task
        self._submit_worker_task = None
        self._stick = None
        _LOGGER.debug("queue stopped")

    async def submit(self, request: PlugwiseRequest) -> PlugwiseResponse | None:
        """Add request to queue and return the received node-response when applicable.

        Raises an error when something fails.
        """
        if request.waiting_for_response:
            raise MessageError(
                f"Cannot send message {request} which is currently waiting for response."
            )

        while True:
            if not request.resend:
                break
            _LOGGER.debug("submit | start (%s) %s", request.retries_left, request)
            if not self._running or self._stick is None:
                raise StickError(
                    f"Cannot send message {request.__class__.__name__} for "
                    f"{request.mac_decoded} because queue manager is stopped"
                )

            await self._add_request_to_queue(request)
            if request.no_response:
                return None

            try:
                response: PlugwiseResponse = await request.response_future()
                return response
            except (NodeTimeout, StickTimeout) as exc:
                if isinstance(request, NodePingRequest):
                    # For ping requests it is expected to receive timeouts, so lower log level
                    _LOGGER.debug(
                        "%s, cancel because timeout is expected for NodePingRequests",
                        exc,
                    )
                    continue
                if request.resend:
                    _LOGGER.debug("%s, retrying", exc)
                    continue
                _LOGGER.debug("%s, cancel request", exc)
                break
            except StickError as exc:
                _LOGGER.error(exc)
                raise StickError(
                    f"No response received for {request.__class__.__name__} "
                    f"to {request.mac_decoded}"
                ) from exc
            except BaseException as exc:
                raise StickError(
                    f"No response received for {request.__class__.__name__} "
                    f"to {request.mac_decoded}"
                ) from exc

        return None

    async def _add_request_to_queue(self, request: PlugwiseRequest) -> None:
        """Add request to send queue."""
        _LOGGER.debug("Add request to queue: %s", request)
        await self._submit_queue.put(request)
        if self._submit_worker_task is None or self._submit_worker_task.done():
            self._submit_worker_task = self._loop.create_task(
                self._send_queue_worker(), name="Send queue worker"
            )

    async def _send_queue_worker(self) -> None:
        """Send messages from queue at the order of priority."""
        _LOGGER.debug("Send_queue_worker started")
        while self._running and self._stick is not None:
            request = await self._submit_queue.get()
            _LOGGER.debug("Sending from send queue %s", request)
            if request.priority == Priority.CANCEL:
                self._submit_queue.task_done()
                return

            qsize = self._submit_queue.qsize()
            if qsize > REPORT_QUEUE_FILLING_UP:
                # When the queue size grows, rate-limit the sending of requests to avoid overloading the network
                await sleep(0.125)
                if qsize > REPORT_QUEUE_FILLING_UP:
                    _LOGGER.warning("Awaiting plugwise responses %d", qsize)

            await self._stick.write_to_stick(request)
            self._submit_queue.task_done()

            _LOGGER.debug("Sent from queue %s", request)
        _LOGGER.debug("Send_queue_worker stopped")