From 642db6d92bd2cd3526dfc30b52f91cd26dcdb262 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Tue, 14 Nov 2023 22:14:37 -0600
Subject: [PATCH] Speed up OTAs (#5720)

---
 esphome/espota2.py | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

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)