diff --git a/esphome/espota2.py b/esphome/espota2.py index 98d6d3a0d9..dbf48a989a 100644 --- a/esphome/espota2.py +++ b/esphome/espota2.py @@ -1,10 +1,13 @@ +from __future__ import annotations + +import gzip import hashlib +import io import logging import random import socket import sys import time -import gzip from esphome.core import EsphomeError from esphome.helpers import is_ip_address, resolve_ip_address @@ -40,6 +43,10 @@ MAGIC_BYTES = [0x6C, 0x26, 0xF7, 0x5C, 0x45] FEATURE_SUPPORTS_COMPRESSION = 0x01 + +UPLOAD_BLOCK_SIZE = 8192 +UPLOAD_BUFFER_SIZE = UPLOAD_BLOCK_SIZE * 8 + _LOGGER = logging.getLogger(__name__) @@ -184,7 +191,9 @@ def send_check(sock, data, msg): raise OTAError(f"Error sending {msg}: {err}") from err -def perform_ota(sock, password, file_handle, filename): +def perform_ota( + sock: socket.socket, password: str, file_handle: io.IOBase, filename: str +) -> None: file_contents = file_handle.read() file_size = len(file_contents) _LOGGER.info("Uploading %s (%s bytes)", filename, file_size) @@ -254,14 +263,16 @@ def perform_ota(sock, password, file_handle, filename): sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 0) # Limit send buffer (usually around 100kB) in order to have progress bar # show the actual progress - sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 8192) + + sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, UPLOAD_BUFFER_SIZE) # Set higher timeout during upload - sock.settimeout(20.0) + sock.settimeout(30.0) + start_time = time.perf_counter() offset = 0 progress = ProgressBar() while True: - chunk = upload_contents[offset : offset + 1024] + chunk = upload_contents[offset : offset + UPLOAD_BLOCK_SIZE] if not chunk: break offset += len(chunk) @@ -277,8 +288,9 @@ def perform_ota(sock, password, file_handle, filename): # Enable nodelay for last checks sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + duration = time.perf_counter() - start_time - _LOGGER.info("Waiting for result...") + _LOGGER.info("Upload took %.2f seconds, waiting for result...", duration) receive_exactly(sock, 1, "receive OK", RESPONSE_RECEIVE_OK) receive_exactly(sock, 1, "Update end", RESPONSE_UPDATE_END_OK)