Merge branch 'ota_speed_up' into integration

This commit is contained in:
J. Nick Koston 2023-11-09 16:59:58 -06:00
commit 905ea480c4
No known key found for this signature in database

View file

@ -1,10 +1,13 @@
from __future__ import annotations
import gzip
import hashlib import hashlib
import io
import logging import logging
import random import random
import socket import socket
import sys import sys
import time import time
import gzip
from esphome.core import EsphomeError from esphome.core import EsphomeError
from esphome.helpers import is_ip_address, resolve_ip_address 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 FEATURE_SUPPORTS_COMPRESSION = 0x01
UPLOAD_BLOCK_SIZE = 8192
UPLOAD_BUFFER_SIZE = UPLOAD_BLOCK_SIZE * 8
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@ -184,7 +191,9 @@ def send_check(sock, data, msg):
raise OTAError(f"Error sending {msg}: {err}") from err 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_contents = file_handle.read()
file_size = len(file_contents) file_size = len(file_contents)
_LOGGER.info("Uploading %s (%s bytes)", filename, file_size) _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) sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 0)
# Limit send buffer (usually around 100kB) in order to have progress bar # Limit send buffer (usually around 100kB) in order to have progress bar
# show the actual progress # 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 # Set higher timeout during upload
sock.settimeout(20.0) sock.settimeout(30.0)
start_time = time.perf_counter()
offset = 0 offset = 0
progress = ProgressBar() progress = ProgressBar()
while True: while True:
chunk = upload_contents[offset : offset + 1024] chunk = upload_contents[offset : offset + UPLOAD_BLOCK_SIZE]
if not chunk: if not chunk:
break break
offset += len(chunk) offset += len(chunk)
@ -277,8 +288,9 @@ def perform_ota(sock, password, file_handle, filename):
# Enable nodelay for last checks # Enable nodelay for last checks
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) 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, "receive OK", RESPONSE_RECEIVE_OK)
receive_exactly(sock, 1, "Update end", RESPONSE_UPDATE_END_OK) receive_exactly(sock, 1, "Update end", RESPONSE_UPDATE_END_OK)