Treat a UDP send failure on the last chunk like a lost final ack

Surface a device error code waiting on the channel, and do not retry once
the whole image was queued, since the device may already be committing.

esphome/espota2.py

@@ -136,7 +136,7 @@ UDP_WINDOW = 4
 # Probe this long before falling back to TCP; the device waits 70 s
 UDP_PROBE_TIMEOUT = 5.0
 UDP_PROBE_INTERVAL = 0.1
-# Resend timeout bounds around srtt + 4 * rttvar; loss never grows it
+# Resend timeout bounds around srtt + 4 * rttvar; random loss never grows it
 UDP_RESEND_INITIAL = 0.5
 UDP_RESEND_MIN = 0.2
 UDP_RESEND_MAX = 2.0
@@ -604,8 +604,9 @@ class UdpChannel:
     """Byte stream over numbered UDP datagrams, standing in for the socket.
 
     The device acks every datagram on arrival and says how many it has room
-    for; each is resent after TCP's srtt + 4 * rttvar, without backoff, so loss
-    slows the transfer instead of ending it.
+    for; each is resent after TCP's srtt + 4 * rttvar, so loss slows the
+    transfer instead of ending it. It backs off only while acks come back
+    untimed, i.e. the link is slower than the timer, never on random loss.
     """
 
     def __init__(self, sock: socket.socket, token: bytes) -> None:
@@ -1173,6 +1174,16 @@ def perform_ota(
                             "No device error behind the send failure: %s", probe_err
                         )
                     raise OTANetworkError(f"sending data: {err}") from err
+                except OTANetworkError as err:
+                    if udp is None:
+                        raise
+                    # Over UDP: surface a device error code if one arrived, and
+                    # treat a failure on the last chunk like a lost final ack
+                    if udp.pending():
+                        receive_exactly(sock, 1, "data", [])
+                    if offset < upload_size:
+                        raise
+                    raise _committed_error(err) from err
 
                 # The UDP channel acks every message itself, so anything the device
                 # sends mid upload is an error code

tests/unit_tests/test_espota2.py

@@ -1986,3 +1986,49 @@ def test_udp_channel_backs_off_when_acks_come_back_untimed() -> None:
     assert rtos[-1] == espota2.UDP_RESEND_MAX
     channel._sample_rtt(0.1)
     assert channel._rto() == pytest.approx(0.3)  # srtt 0.1 + 4 * rttvar 0.05
+
+
+@pytest.mark.usefixtures("mock_time")
+@pytest.mark.parametrize(
+    ("size", "pending", "match", "retryable"),
+    [
+        (100, False, "may have already committed", False),  # one, final chunk
+        (20000, False, "no progress", True),  # fails on an earlier chunk
+        (20000, True, "Writing OTA data", False),  # device error waiting
+    ],
+)
+def test_perform_ota_over_udp_send_failure(
+    mock_socket: Mock, size: int, pending: bool, match: str, retryable: bool
+) -> None:
+    """A UDP send failure surfaces a waiting device error, and on the last
+    chunk is not retried, like a lost final chunk ack over TCP."""
+    mock_socket.recv.side_effect = _udp_handshake()
+    udp = Mock()
+    udp.pending.return_value = pending
+    udp.recv.side_effect = [
+        bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]),
+        bytes([espota2.RESPONSE_BIN_MD5_OK]),
+        bytes([espota2.RESPONSE_ERROR_WRITING_FLASH]),
+    ]
+    # ota type, size and MD5 go through; the first data chunk fails
+    udp.sendall.side_effect = [None] * 3 + [espota2.OTANetworkError("no progress")]
+    with (
+        patch.object(espota2, "_start_udp", return_value=udp),
+        pytest.raises(espota2.OTAError, match=match) as exc,
+    ):
+        espota2.perform_ota(
+            mock_socket, None, io.BytesIO(bytes(size)), "test.bin", prefer_udp=True
+        )
+    assert isinstance(exc.value, espota2.OTANetworkError) is retryable
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_tcp_send_network_error_unchanged(mock_socket: Mock) -> None:
+    """Over TCP an OTANetworkError from a send stays retryable, even on the
+    last chunk."""
+    mock_socket.recv.side_effect = _no_auth_handshake(espota2.OTA_VERSION_2_0)
+    mock_socket.sendall.side_effect = [None] * 4 + [
+        espota2.OTANetworkError("sending frame")
+    ]
+    with pytest.raises(espota2.OTANetworkError, match="sending frame"):
+        espota2.perform_ota(mock_socket, None, io.BytesIO(b"x" * 100), "test.bin")