From d5a19064bac1d55f0a579f8cad6d275fdc7da5ab Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 12 Aug 2026 20:03:04 -0500 Subject: [PATCH] Close the final chunk ack retry window and surface pending device errors --- esphome/espota2.py | 55 +++++++++++++++++------- tests/unit_tests/test_espota2.py | 72 ++++++++++++++++++++++++++++---- 2 files changed, 104 insertions(+), 23 deletions(-) diff --git a/esphome/espota2.py b/esphome/espota2.py index 5d0d28c22e..979e97947a 100644 --- a/esphome/espota2.py +++ b/esphome/espota2.py @@ -79,7 +79,9 @@ UPLOAD_BUFFER_SIZE = UPLOAD_BLOCK_SIZE * 8 # Flaky Wi-Fi links often drop the first OTA attempt, and the device may need time # to clean up a half-open connection (its handshake watchdog runs at 20s) before it # accepts a new one, so wait between attempts instead of failing the upload outright. -MAX_UPLOAD_ATTEMPTS = 3 +# Every resolved address is tried once, and this many extra attempts are shared +# across the addresses on top of that. +EXTRA_UPLOAD_ATTEMPTS = 2 UPLOAD_RETRY_DELAY = 5.0 _LOGGER = logging.getLogger(__name__) @@ -181,6 +183,19 @@ class OTANetworkError(OTAError): """Network-level OTA failure (timeout, reset, closed connection); retrying may succeed.""" +def _committed_error(err: OTANetworkError) -> OTAError: + """Wrap a network failure that happened once the device had the full image. + + Past that point the device commits and reboots on its own, so the failure + must not be retried; a re-upload could flash a device that already updated. + """ + return OTAError( + f"{err} (the device may have already committed the update and " + f"be rebooting; check whether it comes back with the new " + f"firmware before uploading again)" + ) + + def recv_decode( sock: socket.socket, amount: int, decode: bool = True ) -> bytes | list[int]: @@ -473,11 +488,25 @@ def perform_ota( try: sock.sendall(chunk) - if version >= OTA_VERSION_2_0: - receive_exactly(sock, 1, "chunk result", RESPONSE_CHUNK_OK) except OSError as err: + # A send failure can hide an error byte the device reported + # just before dropping the connection; surface that as the + # real, non-retryable cause when it is available + with contextlib.suppress(OSError, OTANetworkError): + sock.settimeout(1.0) + check_error(recv_decode(sock, 1), None) raise OTANetworkError(f"sending data: {err}") from err + if version >= OTA_VERSION_2_0: + try: + receive_exactly(sock, 1, "chunk result", RESPONSE_CHUNK_OK) + except OTANetworkError as err: + if offset < upload_size: + raise + # The device already had the complete image when this ack + # was lost, so it may be committing; do not retry + raise _committed_error(err) from err + progress.update(offset / upload_size) except OTAError: # Terminate the progress bar line before the error is logged @@ -499,11 +528,7 @@ def perform_ota( receive_exactly(sock, 1, "update receive result", RESPONSE_RECEIVE_OK) receive_exactly(sock, 1, "update end result", RESPONSE_UPDATE_END_OK) except OTANetworkError as err: - raise OTAError( - f"{err} (the device may have already committed the update and " - f"be rebooting; check whether it comes back with the new " - f"firmware before uploading again)" - ) from err + raise _committed_error(err) from err try: send_check(sock, RESPONSE_OK, "end acknowledgement") @@ -551,13 +576,13 @@ def run_ota_impl_( _LOGGER.error("No addresses to connect to for %s", remote_host) return 1, None - # Every address is tried at least once and the budget grants - # MAX_UPLOAD_ATTEMPTS - 1 extra retries, cycling through the addresses. - # Wait before an attempt when the previous one actually reached the - # device, or when revisiting an address, so a flaky link can recover and - # the device can clean up a half-open connection (its handshake watchdog - # runs at 20s); moving on to the next address family stays immediate. - total_attempts = len(res) + MAX_UPLOAD_ATTEMPTS - 1 + # Every address is tried at least once and EXTRA_UPLOAD_ATTEMPTS retries + # are shared across the addresses, cycling through them. Wait before an + # attempt when the previous one actually reached the device, or when + # revisiting an address, so a flaky link can recover and the device can + # clean up a half-open connection (its handshake watchdog runs at 20s); + # moving on to the next address family stays immediate. + total_attempts = len(res) + EXTRA_UPLOAD_ATTEMPTS last_error = "" reached_device = False for attempt in range(total_attempts): diff --git a/tests/unit_tests/test_espota2.py b/tests/unit_tests/test_espota2.py index 3848915778..7819c48244 100644 --- a/tests/unit_tests/test_espota2.py +++ b/tests/unit_tests/test_espota2.py @@ -593,7 +593,10 @@ def _no_auth_handshake(version: int) -> list[bytes]: @pytest.mark.usefixtures("mock_time") def test_perform_ota_chunk_send_error(mock_socket: Mock, mock_file: io.BytesIO) -> None: """Test OTA raises the retryable OTANetworkError when sending a chunk fails.""" - mock_socket.recv.side_effect = _no_auth_handshake(espota2.OTA_VERSION_2_0) + mock_socket.recv.side_effect = [ + *_no_auth_handshake(espota2.OTA_VERSION_2_0), + OSError("Connection reset"), # Probe for a pending error byte fails too + ] # Sends before the data phase: magic bytes, features, binary size, MD5; # fail on the fifth sendall, the first firmware chunk mock_socket.sendall.side_effect = [None] * 4 + [OSError("Broken pipe")] @@ -602,6 +605,59 @@ def test_perform_ota_chunk_send_error(mock_socket: Mock, mock_file: io.BytesIO) espota2.perform_ota(mock_socket, None, mock_file, "test.bin") +@pytest.mark.usefixtures("mock_time") +def test_perform_ota_chunk_send_error_surfaces_device_error( + mock_socket: Mock, mock_file: io.BytesIO +) -> None: + """Test a device error byte pending behind a send failure becomes the cause.""" + mock_socket.recv.side_effect = [ + *_no_auth_handshake(espota2.OTA_VERSION_1_0), + bytes([espota2.RESPONSE_ERROR_WRITING_FLASH]), # Reason the device closed + ] + mock_socket.sendall.side_effect = [None] * 4 + [OSError("Broken pipe")] + + with pytest.raises( + espota2.OTAError, match="Writing OTA data to flash memory failed" + ) as exc: + espota2.perform_ota(mock_socket, None, mock_file, "test.bin") + + # The device-reported error is not retryable + assert not isinstance(exc.value, espota2.OTANetworkError) + + +@pytest.mark.usefixtures("mock_time") +def test_perform_ota_final_chunk_ack_failure_not_retryable( + mock_socket: Mock, mock_file: io.BytesIO +) -> None: + """Test a lost ack for the final chunk is not retried.""" + mock_socket.recv.side_effect = [ + *_no_auth_handshake(espota2.OTA_VERSION_2_0), + OSError("Connection reset"), # Ack for the only (final) chunk is lost + ] + + with pytest.raises(espota2.OTAError, match="receiving chunk result") as exc: + espota2.perform_ota(mock_socket, None, mock_file, "test.bin") + + # The device already had the whole image, so it may be committing + assert not isinstance(exc.value, espota2.OTANetworkError) + + +@pytest.mark.usefixtures("mock_time") +def test_perform_ota_intermediate_chunk_ack_failure_retryable( + mock_socket: Mock, +) -> None: + """Test a lost ack for a non-final chunk stays retryable.""" + # Two chunks: the firmware is larger than one upload block + big_file = io.BytesIO(b"x" * (espota2.UPLOAD_BLOCK_SIZE + 1)) + mock_socket.recv.side_effect = [ + *_no_auth_handshake(espota2.OTA_VERSION_2_0), + OSError("Connection reset"), # Ack for the first of two chunks is lost + ] + + with pytest.raises(espota2.OTANetworkError, match="receiving chunk result"): + espota2.perform_ota(mock_socket, None, big_file, "test.bin") + + @pytest.mark.usefixtures("mock_time") def test_perform_ota_post_commit_failure_not_retryable( mock_socket: Mock, mock_file: io.BytesIO @@ -688,9 +744,9 @@ def test_run_ota_impl_connection_failed( assert result_host is None # A single address gets the whole attempt budget, with a delay before # each revisit - assert mock_socket.connect.call_count == espota2.MAX_UPLOAD_ATTEMPTS - assert mock_socket.close.call_count == espota2.MAX_UPLOAD_ATTEMPTS - assert mock_sleep.call_count == espota2.MAX_UPLOAD_ATTEMPTS - 1 + assert mock_socket.connect.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + 1 + assert mock_socket.close.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + 1 + assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS mock_sleep.assert_called_with(espota2.UPLOAD_RETRY_DELAY) @@ -745,8 +801,8 @@ def test_run_ota_impl_network_error_exhausts_attempts( assert result_code == 1 assert result_host is None - assert mock_perform_ota.call_count == espota2.MAX_UPLOAD_ATTEMPTS - assert mock_sleep.call_count == espota2.MAX_UPLOAD_ATTEMPTS - 1 + assert mock_perform_ota.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + 1 + assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS @pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip_dual") @@ -762,8 +818,8 @@ def test_run_ota_impl_multiple_addresses_cycle( assert result_code == 1 assert result_host is None - # Each address is visited once, then the two spare attempts cycle back - # through them; the budget is shared, not per address + # Each address is visited once, then the EXTRA_UPLOAD_ATTEMPTS spare + # attempts cycle back through them; the budget is shared, not per address assert mock_socket.connect.call_args_list == [ call(DUAL_STACK_SA6), call(DUAL_STACK_SA4),