| # Copyright 2022 The Pigweed Authors |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); you may not |
| # use this file except in compliance with the License. You may obtain a copy of |
| # the License at |
| # |
| # https://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT |
| # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the |
| # License for the specific language governing permissions and limitations under |
| # the License. |
| """Classes for read and write transfers.""" |
| |
| import abc |
| import asyncio |
| from dataclasses import dataclass |
| import enum |
| import logging |
| import math |
| import threading |
| from typing import Any, Callable, Optional |
| |
| from pw_status import Status |
| try: |
| from pw_transfer import transfer_pb2 |
| except ImportError: |
| # For the bazel build, which puts generated protos in a different location. |
| from pigweed.pw_transfer import transfer_pb2 # type: ignore |
| |
| _LOG = logging.getLogger(__package__) |
| |
| |
| @dataclass(frozen=True) |
| class ProgressStats: |
| bytes_sent: int |
| bytes_confirmed_received: int |
| total_size_bytes: Optional[int] |
| |
| def percent_received(self) -> float: |
| if self.total_size_bytes is None: |
| return math.nan |
| |
| return self.bytes_confirmed_received / self.total_size_bytes * 100 |
| |
| def __str__(self) -> str: |
| total = str( |
| self.total_size_bytes) if self.total_size_bytes else 'unknown' |
| return (f'{self.percent_received():5.1f}% ({self.bytes_sent} B sent, ' |
| f'{self.bytes_confirmed_received} B received of {total} B)') |
| |
| |
| ProgressCallback = Callable[[ProgressStats], Any] |
| |
| |
| class _Timer: |
| """A timer which invokes a callback after a certain timeout.""" |
| def __init__(self, timeout_s: float, callback: Callable[[], Any]): |
| self.timeout_s = timeout_s |
| self._callback = callback |
| self._task: Optional[asyncio.Task[Any]] = None |
| |
| def start(self, timeout_s: float = None) -> None: |
| """Starts a new timer. |
| |
| If a timer is already running, it is stopped and a new timer started. |
| This can be used to implement watchdog-like behavior, where a callback |
| is invoked after some time without a kick. |
| """ |
| self.stop() |
| timeout_s = self.timeout_s if timeout_s is None else timeout_s |
| self._task = asyncio.create_task(self._run(timeout_s)) |
| |
| def stop(self) -> None: |
| """Terminates a running timer.""" |
| if self._task is not None: |
| self._task.cancel() |
| self._task = None |
| |
| async def _run(self, timeout_s: float) -> None: |
| await asyncio.sleep(timeout_s) |
| self._task = None |
| self._callback() |
| |
| |
| class Transfer(abc.ABC): |
| """A client-side data transfer through a Manager. |
| |
| Subclasses are responsible for implementing all of the logic for their type |
| of transfer, receiving messages from the server and sending the appropriate |
| messages in response. |
| """ |
| class _State(enum.Enum): |
| # Transfer is starting. The server and client are performing an initial |
| # handshake and negotiating protocol and feature flags. |
| INITIATING = 0 |
| |
| # Waiting for the other end to send a chunk. |
| WAITING = 1 |
| |
| # Transmitting a window of data to a receiver. |
| TRANSMITTING = 2 |
| |
| # Recovering after one or more chunks was dropped in an active transfer. |
| RECOVERY = 3 |
| |
| # Transfer has completed locally and is waiting for the peer to |
| # acknowledge its final status. Only entered by the terminating side of |
| # the transfer. |
| # |
| # The context remains in a TERMINATING state until it receives an |
| # acknowledgement from the peer or times out. |
| TERMINATING = 4 |
| |
| # A transfer has fully completed. |
| COMPLETE = 5 |
| |
| def __init__(self, |
| session_id: int, |
| send_chunk: Callable[[transfer_pb2.Chunk], None], |
| end_transfer: Callable[['Transfer'], None], |
| response_timeout_s: float, |
| initial_response_timeout_s: float, |
| max_retries: int, |
| progress_callback: ProgressCallback = None): |
| self.id = session_id |
| self.status = Status.OK |
| self.done = threading.Event() |
| |
| self._send_chunk = send_chunk |
| self._end_transfer = end_transfer |
| |
| # TODO(b/243679452): Only the legacy protocol begins in a WAITING state. |
| # Newer protocols should start as INITIATING. |
| self._state = Transfer._State.WAITING |
| |
| self._retries = 0 |
| self._max_retries = max_retries |
| self._response_timer = _Timer(response_timeout_s, self._on_timeout) |
| self._initial_response_timeout_s = initial_response_timeout_s |
| |
| self._progress_callback = progress_callback |
| |
| async def begin(self) -> None: |
| """Sends the initial chunk of the transfer.""" |
| self._send_chunk(self._initial_chunk()) |
| self._response_timer.start(self._initial_response_timeout_s) |
| |
| @property |
| @abc.abstractmethod |
| def data(self) -> bytes: |
| """Returns the data read or written in this transfer.""" |
| |
| @abc.abstractmethod |
| def _initial_chunk(self) -> transfer_pb2.Chunk: |
| """Returns the initial chunk to notify the sever of the transfer.""" |
| |
| async def handle_chunk(self, chunk: transfer_pb2.Chunk) -> None: |
| """Processes an incoming chunk from the server. |
| |
| Handles terminating chunks (i.e. those with a status) and forwards |
| non-terminating chunks to handle_data_chunk. |
| """ |
| self._response_timer.stop() |
| self._retries = 0 # Received data from service, so reset the retries. |
| |
| _LOG.debug('Received chunk\n%s', str(chunk).rstrip()) |
| |
| # Status chunks are only used to terminate a transfer. They do not |
| # contain any data that requires processing. |
| if chunk.HasField('status'): |
| self.finish(Status(chunk.status)) |
| return |
| |
| await self._handle_data_chunk(chunk) |
| |
| # Start the timeout for the server to send a chunk in response. |
| self._response_timer.start() |
| |
| @abc.abstractmethod |
| async def _handle_data_chunk(self, chunk: transfer_pb2.Chunk) -> None: |
| """Handles a chunk that contains or requests data.""" |
| |
| @abc.abstractmethod |
| def _retry_after_timeout(self) -> None: |
| """Retries after a timeout occurs.""" |
| |
| def _on_timeout(self) -> None: |
| """Handles a timeout while waiting for a chunk.""" |
| if self._state is Transfer._State.COMPLETE: |
| return |
| |
| self._retries += 1 |
| if self._retries > self._max_retries: |
| self.finish(Status.DEADLINE_EXCEEDED) |
| return |
| |
| _LOG.debug('Received no responses for %.3fs; retrying %d/%d', |
| self._response_timer.timeout_s, self._retries, |
| self._max_retries) |
| self._retry_after_timeout() |
| self._response_timer.start() |
| |
| def finish(self, status: Status, skip_callback: bool = False) -> None: |
| """Ends the transfer with the specified status.""" |
| self._response_timer.stop() |
| self.status = status |
| |
| if status.ok(): |
| total_size = len(self.data) |
| self._update_progress(total_size, total_size, total_size) |
| |
| if not skip_callback: |
| self._end_transfer(self) |
| |
| # Set done last so that the transfer has been fully cleaned up. |
| self._state = Transfer._State.COMPLETE |
| self.done.set() |
| |
| def _update_progress(self, bytes_sent: int, bytes_confirmed_received: int, |
| total_size_bytes: Optional[int]) -> None: |
| """Invokes the provided progress callback, if any, with the progress.""" |
| |
| stats = ProgressStats(bytes_sent, bytes_confirmed_received, |
| total_size_bytes) |
| _LOG.debug('Transfer %d progress: %s', self.id, stats) |
| |
| if self._progress_callback: |
| self._progress_callback(stats) |
| |
| def _send_error(self, error: Status) -> None: |
| """Sends an error chunk to the server and finishes the transfer.""" |
| self._send_chunk( |
| transfer_pb2.Chunk(transfer_id=self.id, |
| status=error.value, |
| type=transfer_pb2.Chunk.Type.COMPLETION)) |
| self.finish(error) |
| |
| |
| class WriteTransfer(Transfer): |
| """A client -> server write transfer.""" |
| def __init__( |
| self, |
| session_id: int, |
| data: bytes, |
| send_chunk: Callable[[transfer_pb2.Chunk], None], |
| end_transfer: Callable[[Transfer], None], |
| response_timeout_s: float, |
| initial_response_timeout_s: float, |
| max_retries: int, |
| progress_callback: ProgressCallback = None, |
| ): |
| super().__init__(session_id, send_chunk, end_transfer, |
| response_timeout_s, initial_response_timeout_s, |
| max_retries, progress_callback) |
| self._data = data |
| |
| self._offset = 0 |
| self._window_end_offset = 0 |
| self._max_chunk_size = 0 |
| self._chunk_delay_us: Optional[int] = None |
| |
| # The window ID increments for each parameters update. |
| self._window_id = 0 |
| |
| self._bytes_confirmed_received = 0 |
| self._last_chunk = self._initial_chunk() |
| |
| @property |
| def data(self) -> bytes: |
| return self._data |
| |
| def _initial_chunk(self) -> transfer_pb2.Chunk: |
| # TODO(b/243679452): In protocol v2, session_id should be assigned by |
| # the server during an initial handshake. |
| return transfer_pb2.Chunk(transfer_id=self.id, |
| resource_id=self.id, |
| type=transfer_pb2.Chunk.Type.START) |
| |
| async def _handle_data_chunk(self, chunk: transfer_pb2.Chunk) -> None: |
| """Processes an incoming chunk from the server. |
| |
| In a write transfer, the server only sends transfer parameter updates |
| to the client. When a message is received, update local parameters and |
| send data accordingly. |
| """ |
| |
| if self._state is Transfer._State.TERMINATING: |
| # TODO(b/243679452): Handle termination in protocol v2. |
| return |
| |
| if self._state is Transfer._State.TRANSMITTING: |
| self._state = Transfer._State.WAITING |
| |
| assert self._state is Transfer._State.WAITING |
| |
| if not self._handle_parameters_update(chunk): |
| return |
| |
| self._bytes_confirmed_received = chunk.offset |
| self._state = Transfer._State.TRANSMITTING |
| |
| self._window_id += 1 |
| asyncio.create_task(self._transmit_next_chunk(self._window_id)) |
| |
| async def _transmit_next_chunk(self, |
| window_id: int, |
| timeout_us: int = None) -> None: |
| """Transmits a single data chunk to the server. |
| |
| If the chunk completes the active window, returns to a WAITING state. |
| Otherwise, schedules another transmission for the next chunk. |
| """ |
| if timeout_us is not None: |
| await asyncio.sleep(timeout_us / 1e6) |
| |
| if self._state is not Transfer._State.TRANSMITTING: |
| return |
| |
| if window_id != self._window_id: |
| _LOG.debug('Transfer %d: Skipping stale window', self.id) |
| return |
| |
| chunk = self._next_chunk() |
| self._offset += len(chunk.data) |
| |
| sent_requested_bytes = self._offset == self._window_end_offset |
| |
| self._send_chunk(chunk) |
| self._update_progress(self._offset, self._bytes_confirmed_received, |
| len(self.data)) |
| |
| if sent_requested_bytes: |
| self._state = Transfer._State.WAITING |
| else: |
| asyncio.create_task( |
| self._transmit_next_chunk(window_id, |
| timeout_us=self._chunk_delay_us)) |
| |
| self._last_chunk = chunk |
| |
| def _handle_parameters_update(self, chunk: transfer_pb2.Chunk) -> bool: |
| """Updates transfer state based on a transfer parameters update.""" |
| |
| retransmit = True |
| if chunk.HasField('type'): |
| retransmit = (chunk.type |
| == transfer_pb2.Chunk.Type.PARAMETERS_RETRANSMIT |
| or chunk.type == transfer_pb2.Chunk.Type.START) |
| |
| if chunk.offset > len(self.data): |
| # Bad offset; terminate the transfer. |
| _LOG.error( |
| 'Transfer %d: server requested invalid offset %d (size %d)', |
| self.id, chunk.offset, len(self.data)) |
| |
| self._send_error(Status.OUT_OF_RANGE) |
| return False |
| |
| if chunk.pending_bytes == 0: |
| _LOG.error( |
| 'Transfer %d: service requested 0 bytes (invalid); aborting', |
| self.id) |
| self._send_error(Status.INTERNAL) |
| return False |
| |
| if retransmit: |
| # Check whether the client has sent a previous data offset, which |
| # indicates that some chunks were lost in transmission. |
| if chunk.offset < self._offset: |
| _LOG.debug('Write transfer %d rolling back: offset %d from %d', |
| self.id, chunk.offset, self._offset) |
| |
| self._offset = chunk.offset |
| |
| # Retransmit is the default behavior for older versions of the |
| # transfer protocol. The window_end_offset field is not guaranteed |
| # to be set in these version, so it must be calculated. |
| max_bytes_to_send = min(chunk.pending_bytes, |
| len(self.data) - self._offset) |
| self._window_end_offset = self._offset + max_bytes_to_send |
| else: |
| assert chunk.type == transfer_pb2.Chunk.Type.PARAMETERS_CONTINUE |
| |
| # Extend the window to the new end offset specified by the server. |
| self._window_end_offset = min(chunk.window_end_offset, |
| len(self.data)) |
| |
| if chunk.HasField('max_chunk_size_bytes'): |
| self._max_chunk_size = chunk.max_chunk_size_bytes |
| |
| if chunk.HasField('min_delay_microseconds'): |
| self._chunk_delay_us = chunk.min_delay_microseconds |
| |
| return True |
| |
| def _retry_after_timeout(self) -> None: |
| if self._state is Transfer._State.WAITING: |
| self._send_chunk(self._last_chunk) |
| |
| # TODO(b/243679452): Handle other states in protocol v2. |
| |
| def _next_chunk(self) -> transfer_pb2.Chunk: |
| """Returns the next Chunk message to send in the data transfer.""" |
| chunk = transfer_pb2.Chunk(transfer_id=self.id, |
| offset=self._offset, |
| type=transfer_pb2.Chunk.Type.DATA) |
| max_bytes_in_chunk = min(self._max_chunk_size, |
| self._window_end_offset - self._offset) |
| |
| chunk.data = self.data[self._offset:self._offset + max_bytes_in_chunk] |
| |
| # Mark the final chunk of the transfer. |
| if len(self.data) - self._offset <= max_bytes_in_chunk: |
| chunk.remaining_bytes = 0 |
| |
| return chunk |
| |
| |
| class ReadTransfer(Transfer): |
| """A client <- server read transfer. |
| |
| Although Python can effectively handle an unlimited transfer window, this |
| client sets a conservative window and chunk size to avoid overloading the |
| device. These are configurable in the constructor. |
| """ |
| |
| # The fractional position within a window at which a receive transfer should |
| # extend its window size to minimize the amount of time the transmitter |
| # spends blocked. |
| # |
| # For example, a divisor of 2 will extend the window when half of the |
| # requested data has been received, a divisor of three will extend at a |
| # third of the window, and so on. |
| EXTEND_WINDOW_DIVISOR = 2 |
| |
| def __init__( # pylint: disable=too-many-arguments |
| self, |
| session_id: int, |
| send_chunk: Callable[[transfer_pb2.Chunk], None], |
| end_transfer: Callable[[Transfer], None], |
| response_timeout_s: float, |
| initial_response_timeout_s: float, |
| max_retries: int, |
| max_bytes_to_receive: int = 8192, |
| max_chunk_size: int = 1024, |
| chunk_delay_us: int = None, |
| progress_callback: ProgressCallback = None): |
| super().__init__(session_id, send_chunk, end_transfer, |
| response_timeout_s, initial_response_timeout_s, |
| max_retries, progress_callback) |
| self._max_bytes_to_receive = max_bytes_to_receive |
| self._max_chunk_size = max_chunk_size |
| self._chunk_delay_us = chunk_delay_us |
| |
| self._remaining_transfer_size: Optional[int] = None |
| self._data = bytearray() |
| self._offset = 0 |
| self._pending_bytes = max_bytes_to_receive |
| self._window_end_offset = max_bytes_to_receive |
| self._last_chunk_offset: Optional[int] = None |
| |
| @property |
| def data(self) -> bytes: |
| """Returns an immutable copy of the data that has been read.""" |
| return bytes(self._data) |
| |
| def _initial_chunk(self) -> transfer_pb2.Chunk: |
| return self._transfer_parameters(transfer_pb2.Chunk.Type.START) |
| |
| async def _handle_data_chunk(self, chunk: transfer_pb2.Chunk) -> None: |
| """Processes an incoming chunk from the server. |
| |
| In a read transfer, the client receives data chunks from the server. |
| Once all pending data is received, the transfer parameters are updated. |
| """ |
| |
| if self._state is Transfer._State.TERMINATING: |
| # TODO(b/243679452): Handle terminating states in protocol v2. |
| return |
| |
| if self._state is Transfer._State.RECOVERY: |
| if chunk.offset != self._offset: |
| if self._last_chunk_offset == chunk.offset: |
| _LOG.debug( |
| 'Transfer %d received repeated offset %d: ' |
| 'retry detected, resending transfer parameters', |
| self.id, chunk.offset) |
| self._send_chunk( |
| self._transfer_parameters( |
| transfer_pb2.Chunk.Type.PARAMETERS_RETRANSMIT)) |
| else: |
| _LOG.debug( |
| 'Transfer %d waiting for offset %d, ignoring %d', |
| self.id, self._offset, chunk.offset) |
| self._last_chunk_offset = chunk.offset |
| return |
| |
| _LOG.info( |
| 'Transfer %d received expected offset %d, resuming transfer', |
| self.id, chunk.offset) |
| self._state = Transfer._State.WAITING |
| |
| assert self._state is Transfer._State.WAITING |
| |
| if chunk.offset != self._offset: |
| # Initially, the transfer service only supports in-order transfers. |
| # If data is received out of order, request that the server |
| # retransmit from the previous offset. |
| _LOG.debug( |
| 'Transfer %d expected offset %d, received %d: ' |
| 'entering recovery state', self.id, self._offset, chunk.offset) |
| self._state = Transfer._State.RECOVERY |
| |
| self._send_chunk( |
| self._transfer_parameters( |
| transfer_pb2.Chunk.Type.PARAMETERS_RETRANSMIT)) |
| return |
| |
| self._data += chunk.data |
| self._pending_bytes -= len(chunk.data) |
| self._offset += len(chunk.data) |
| |
| # Update the last offset seen so that retries can be detected. |
| self._last_chunk_offset = chunk.offset |
| |
| if chunk.HasField('remaining_bytes'): |
| if chunk.remaining_bytes == 0: |
| # No more data to read. Acknowledge receipt and finish. |
| self._send_chunk( |
| transfer_pb2.Chunk( |
| transfer_id=self.id, |
| status=Status.OK.value, |
| type=transfer_pb2.Chunk.Type.COMPLETION)) |
| self.finish(Status.OK) |
| return |
| |
| # The server may indicate if the amount of remaining data is known. |
| self._remaining_transfer_size = chunk.remaining_bytes |
| elif self._remaining_transfer_size is not None: |
| # Update the remaining transfer size, if it is known. |
| self._remaining_transfer_size -= len(chunk.data) |
| |
| # If the transfer size drops to zero, the estimate was inaccurate. |
| if self._remaining_transfer_size <= 0: |
| self._remaining_transfer_size = None |
| |
| total_size = None if self._remaining_transfer_size is None else ( |
| self._remaining_transfer_size + self._offset) |
| self._update_progress(self._offset, self._offset, total_size) |
| |
| if chunk.window_end_offset != 0: |
| if chunk.window_end_offset < self._offset: |
| _LOG.error( |
| 'Transfer %d: transmitter sent invalid earlier end offset ' |
| '%d (receiver offset %d)', self.id, |
| chunk.window_end_offset, self._offset) |
| self._send_error(Status.INTERNAL) |
| return |
| |
| if chunk.window_end_offset > self._window_end_offset: |
| _LOG.error( |
| 'Transfer %d: transmitter sent invalid later end offset ' |
| '%d (receiver end offset %d)', self.id, |
| chunk.window_end_offset, self._window_end_offset) |
| self._send_error(Status.INTERNAL) |
| return |
| |
| self._window_end_offset = chunk.window_end_offset |
| self._pending_bytes -= chunk.window_end_offset - self._offset |
| |
| remaining_window_size = self._window_end_offset - self._offset |
| extend_window = (remaining_window_size <= self._max_bytes_to_receive / |
| ReadTransfer.EXTEND_WINDOW_DIVISOR) |
| |
| if self._pending_bytes == 0: |
| # All pending data was received. Send out a new parameters chunk for |
| # the next block. |
| self._send_chunk( |
| self._transfer_parameters( |
| transfer_pb2.Chunk.Type.PARAMETERS_RETRANSMIT)) |
| elif extend_window: |
| self._send_chunk( |
| self._transfer_parameters( |
| transfer_pb2.Chunk.Type.PARAMETERS_CONTINUE)) |
| |
| def _retry_after_timeout(self) -> None: |
| if (self._state is Transfer._State.WAITING |
| or self._state is Transfer._State.RECOVERY): |
| self._send_chunk( |
| self._transfer_parameters( |
| transfer_pb2.Chunk.Type.PARAMETERS_RETRANSMIT)) |
| |
| # TODO(b/243679452): Handle other states in protocol v2. |
| |
| def _transfer_parameters(self, chunk_type: Any) -> transfer_pb2.Chunk: |
| """Sends an updated transfer parameters chunk to the server.""" |
| |
| self._pending_bytes = self._max_bytes_to_receive |
| self._window_end_offset = self._offset + self._max_bytes_to_receive |
| |
| chunk = transfer_pb2.Chunk(transfer_id=self.id, |
| pending_bytes=self._pending_bytes, |
| window_end_offset=self._window_end_offset, |
| max_chunk_size_bytes=self._max_chunk_size, |
| offset=self._offset, |
| type=chunk_type) |
| |
| if self._chunk_delay_us: |
| chunk.min_delay_microseconds = self._chunk_delay_us |
| |
| return chunk |
