Speed up OTAs (#5720)

This commit is contained in:
J. Nick Koston 2023-11-14 22:14:37 -06:00 committed by GitHub
parent 4aac5a23cd
commit 642db6d92b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -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)