Esempio n. 1
0
class Port(object):
    def __init__(self, tag="data", maxsize=1, name=None, loop=None):
        loop = loop if loop is not None else asyncio.get_event_loop()
        self.loop = loop
        self.name = name if name is not None else str(uuid1())
        self._queue = Queue(maxsize, loop=self.loop)
        self.default_value = None
        self.default_value_set = False
        self.connected = False
        self.belong_to_block = None
        self.data_tag = tag

    def set_default_value(self, value):
        if not isinstance(value, Payload):
            raise Exception("value should be Payload type")
        self.default_value = value
        self.default_value_set = True

    async def get(self):
        if self.default_value_set:
            if self._queue.empty():
                return self.default_value, self.default_value[self.data_tag]

        payload = await self._queue.get()
        return payload, payload[self.data_tag]

    def get_nowait(self):
        if self.default_value_set:
            if self._queue.empty():
                return self.default_value, self.default_value[self.data_tag]

        payload = self._queue.get_nowait()
        return payload, payload[self.data_tag]

    async def put(self, payload, item):
        if self.connected:
            payload[self.data_tag] = item
            await self._queue.put(payload)

    def put_nowait(self, payload, item):
        if self.connected:
            payload[self.data_tag] = item
            self._queue.put_nowait(payload)

    def empty(self):
        return self._queue.empty()

    def full(self):
        return self._queue.full()

    def set_buffer_size(self, maxsize):
        self._queue = Queue(maxsize, loop=self.loop)
class GoogleReportState(iot_base.BaseIoT):
    """Report states to Google.

    Uses a queue to send messages.
    """
    def __init__(self, cloud: Cloud):
        """Initialize Google Report State."""
        super().__init__(cloud)
        self._connect_lock = asyncio.Lock()
        self._to_send = Queue(100)
        self._message_sender_task = None
        # Local code waiting for a response
        self._response_handler: Dict[str, asyncio.Future] = {}
        self.register_on_connect(self._async_on_connect)
        self.register_on_disconnect(self._async_on_disconnect)

        # Register start/stop
        cloud.register_on_stop(self.disconnect)

    @property
    def package_name(self) -> str:
        """Return the package name for logging."""
        return __name__

    @property
    def ws_server_url(self) -> str:
        """Server to connect to."""
        # https -> wss, http -> ws
        return f"ws{self.cloud.google_actions_report_state_url[4:]}/v1"

    async def async_send_message(self, msg):
        """Send a message."""
        msgid = uuid.uuid4().hex

        # Since connect is async, guard against send_message called twice in parallel.
        async with self._connect_lock:
            if self.state == iot_base.STATE_DISCONNECTED:
                self.cloud.run_task(self.connect())
                # Give connect time to start up and change state.
                await asyncio.sleep(0)

        if self._to_send.full():
            discard_msg = self._to_send.get_nowait()
            self._response_handler.pop(discard_msg["msgid"]).set_exception(
                ErrorResponse(ERR_DISCARD_CODE, ERR_DISCARD_MSG))

        fut = self._response_handler[msgid] = asyncio.Future()

        self._to_send.put_nowait({"msgid": msgid, "payload": msg})

        try:
            return await fut
        finally:
            self._response_handler.pop(msgid, None)

    def async_handle_message(self, msg):
        """Handle a message."""
        response_handler = self._response_handler.get(msg["msgid"])

        if response_handler is not None:
            if "error" in msg:
                response_handler.set_exception(
                    ErrorResponse(msg["error"], msg["message"]))
            else:
                response_handler.set_result(msg.get("payload"))
            return

        self._logger.warning("Got unhandled message: %s", msg)

    async def _async_on_connect(self):
        """On Connect handler."""
        self._message_sender_task = self.cloud.run_task(
            self._async_message_sender())

    async def _async_on_disconnect(self):
        """On disconnect handler."""
        self._message_sender_task.cancel()
        self._message_sender_task = None

    async def _async_message_sender(self):
        """Start sending messages."""
        self._logger.debug("Message sender task activated")
        try:
            while True:
                await self.async_send_json_message(await self._to_send.get())
        except asyncio.CancelledError:
            pass
        self._logger.debug("Message sender task shut down")