[ota] Restore lazy flash erase for ESP32 OTA with 64 KiB block erase (#18580)

esphome/components/esphome/ota/ota_esphome.cpp

@@ -398,7 +398,7 @@ void ESPHomeOTAComponent::handle_data_() {
   this->notify_state_(ota::OTA_STARTED, 0.0f, 0);
 #endif
 
-  // begin() may block for a few seconds while it locks flash.
+  // begin() returns quickly; flash sectors are erased incrementally during write().
   error_code = this->backend_->begin(ota_size, ota_type);
   if (error_code != ota::OTA_RESPONSE_OK)
     goto error;  // NOLINT(cppcoreguidelines-avoid-goto)

esphome/components/ota/ota_backend.h

@@ -66,6 +66,19 @@ enum OTAResponseTypes {
  */
 bool version_is_older(const char *candidate, const char *reference);
 
+// 64 KiB flash block; the erase granularity the ESP-IDF backend erases ahead with.
+static constexpr size_t OTA_BLOCK_ERASE_SIZE = 64 * 1024;
+
+/** Target erased watermark for lazy block erase-ahead.
+ *
+ * Rounds the write end offset up to a block boundary, clamped to the partition
+ * size. Platform-independent so the arithmetic is host-testable.
+ */
+constexpr size_t next_erase_end(size_t write_end, size_t partition_size) {
+  const size_t rounded = (write_end + OTA_BLOCK_ERASE_SIZE - 1) & ~(OTA_BLOCK_ERASE_SIZE - 1);
+  return rounded < partition_size ? rounded : partition_size;
+}
+
 enum OTAState {
   OTA_COMPLETED = 0,
   OTA_STARTED,

esphome/components/ota/ota_backend_esp_idf.cpp

@@ -7,7 +7,7 @@
 #include "esphome/core/log.h"
 
 #include <esp_ota_ops.h>
-#include <esp_task_wdt.h>
+#include <sdkconfig.h>
 #include <spi_flash_mmap.h>
 #ifdef USE_OTA_DOWNGRADE_PROTECTION
 #include <esp_app_desc.h>
@@ -60,27 +60,38 @@ OTAResponseTypes IDFOTABackend::begin(size_t image_size, ota::OTAType ota_type)
     return OTA_RESPONSE_ERROR_NO_UPDATE_PARTITION;
   }
 
-  // esp_ota_begin() erases the destination region, which blocks loopTask and
-  // scales with the erase size -- a fixed watchdog overruns on large OTA slots.
-  // An unknown size (0, e.g. web_server uploads) erases the whole partition, so
-  // budget against the bytes actually erased. ~10ms/KiB (conservative
-  // ~100 KiB/s erase) over a 15s floor; panic stays on so a stuck erase still
-  // resets rather than hanging forever.
-  size_t erase_size = image_size;
-  if (erase_size == 0 || erase_size > this->partition_->size) {
-    erase_size = this->partition_->size;
+  // Both lazy-erase paths below replace esp_ota_begin()'s blocking full erase.
+  // Size check replaces the one that erase performed (0 = unknown size,
+  // e.g. web_server uploads).
+  if (image_size != 0 && image_size > this->partition_->size) {
+    return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
   }
-  const uint32_t erase_budget_ms = 15000 + (erase_size >> 10) * 10;
-  watchdog::WatchdogManager watchdog(erase_budget_ms);
-  esp_err_t err = esp_ota_begin(this->partition_, image_size, &this->update_handle_);
+  this->written_ = 0;
+  esp_err_t err;
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
+  this->erased_end_ = 0;
+  // Unlike esp_ota_begin(), esp_ota_resume() does not reject a running app in
+  // ESP_OTA_IMG_PENDING_VERIFY; that state is unreachable here because the app
+  // was marked valid at boot (esp32/hal.cpp) or just above under USE_OTA_ROLLBACK.
+  // erase_size 0 (!= OTA_WITH_SEQUENTIAL_WRITES) means no erase; erase_ahead_() handles it
+  err = esp_ota_resume(this->partition_, 0, 0, &this->update_handle_);
+#if defined(CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE) && ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(5, 5, 0)
+  // esp_ota_begin() does this on IDF 5.5+; esp_ota_resume() does not. Prevents
+  // booting a half-written slot after a crash mid-OTA. Not available on the
+  // 5.3.3/5.4.2 backports, whose esp_ota_begin() did not invalidate either.
+  if (err == ESP_OK) {
+    esp_ota_invalidate_inactive_ota_data_slot();
+  }
+#endif
+#else
+  err = esp_ota_begin(this->partition_, OTA_WITH_SEQUENTIAL_WRITES, &this->update_handle_);
+#endif
 
   if (err != ESP_OK) {
-    ESP_LOGE(TAG, "esp_ota_begin failed (err=0x%X)", err);
+    ESP_LOGE(TAG, "OTA begin failed (err=0x%X)", err);
     esp_ota_abort(this->update_handle_);
     this->update_handle_ = 0;
-    if (err == ESP_ERR_INVALID_SIZE) {
-      return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
-    } else if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
+    if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
       return OTA_RESPONSE_ERROR_WRITING_FLASH;
     } else if (err == ESP_ERR_OTA_PARTITION_CONFLICT) {
       // This error appears with 1 factory and 1 ota partition
@@ -120,6 +131,17 @@ OTAResponseTypes IDFOTABackend::write(uint8_t *data, size_t len) {
   if (!this->is_app_or_bootloader_update_()) {
     return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
   }
+#endif
+  // Overflow can only happen on unknown-size uploads (web_server); known
+  // sizes were rejected in begin().
+  if (this->written_ + len > this->partition_->size) {
+    return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
+  }
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
+  OTAResponseTypes erase_result = this->erase_ahead_(len);
+  if (erase_result != OTA_RESPONSE_OK) {
+    return erase_result;
+  }
 #endif
   esp_err_t err = esp_ota_write(this->update_handle_, data, len);
   this->md5_.add(data, len);
@@ -127,14 +149,40 @@ OTAResponseTypes IDFOTABackend::write(uint8_t *data, size_t len) {
     ESP_LOGE(TAG, "esp_ota_write failed (err=0x%X)", err);
     if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
       return OTA_RESPONSE_ERROR_MAGIC;
+    } else if (err == ESP_ERR_INVALID_SIZE) {
+      // Sequential-writes fallback: IDF's lazy erase reports overflow here
+      return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
     } else if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
       return OTA_RESPONSE_ERROR_WRITING_FLASH;
     }
     return OTA_RESPONSE_ERROR_UNKNOWN;
   }
+  this->written_ += len;
   return OTA_RESPONSE_OK;
 }
 
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
+OTAResponseTypes IDFOTABackend::erase_ahead_(size_t len) {
+  const size_t end = this->written_ + len;
+  if (this->erased_end_ >= end) {
+    return OTA_RESPONSE_OK;
+  }
+  // Round up to a block boundary, clamped to the partition end; IDF splits the
+  // range into 64 KiB block erases where aligned, sector erases elsewhere.
+  const size_t erase_to = next_erase_end(end, this->partition_->size);
+  // A block erase is one uninterruptible flash op (typically ~150 ms, seconds
+  // on aged flash) and the transfer loop may not have fed the WDT for ~1s.
+  watchdog::WatchdogManager watchdog(15000);
+  esp_err_t err = esp_partition_erase_range(this->partition_, this->erased_end_, erase_to - this->erased_end_);
+  if (err != ESP_OK) {
+    ESP_LOGE(TAG, "esp_partition_erase_range failed (err=0x%X)", err);
+    return err == ESP_ERR_INVALID_SIZE ? OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE : OTA_RESPONSE_ERROR_WRITING_FLASH;
+  }
+  this->erased_end_ = erase_to;
+  return OTA_RESPONSE_OK;
+}
+#endif
+
 OTAResponseTypes IDFOTABackend::end() {
   if (this->md5_set_) {
     this->md5_.calculate();
@@ -226,6 +274,10 @@ void IDFOTABackend::abort() {
   // or not an update is in flight.
   esp_ota_abort(this->update_handle_);
   this->update_handle_ = 0;
+  this->written_ = 0;
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
+  this->erased_end_ = 0;
+#endif
 }
 
 }  // namespace esphome::ota

esphome/components/ota/ota_backend_esp_idf.h

@@ -5,8 +5,18 @@
 #include "esphome/components/md5/md5.h"
 #include "esphome/core/defines.h"
 
+#include <esp_idf_version.h>
 #include <esp_ota_ops.h>
 
+// esp_ota_resume() (IDF 5.4.2+, backported to 5.3.3) provides a no-erase OTA
+// handle, letting write() block-erase 64 KiB ahead of the write cursor
+// (~4x faster than the per-sector lazy erase of OTA_WITH_SEQUENTIAL_WRITES,
+// used as fallback on older IDF).
+#if ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(5, 4, 2) || \
+    (ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(5, 3, 3) && ESP_IDF_VERSION < ESP_IDF_VERSION_VAL(5, 4, 0))
+#define USE_OTA_BLOCK_ERASE_AHEAD
+#endif
+
 namespace esphome::ota {
 
 #ifdef USE_OTA_PARTITIONS
@@ -54,6 +64,9 @@ class IDFOTABackend final {
 #endif
 
  private:
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
+  OTAResponseTypes erase_ahead_(size_t len);
+#endif
 #ifdef USE_OTA_SIGNED_VERIFICATION_MULTI_KEY
   // Accept an image signed by any key the running app trusts (up to 3 blocks),
   // so rotation and backup keys work. Fails closed. Covers app and bootloader.
@@ -62,7 +75,11 @@ class IDFOTABackend final {
   // Keep md5_ first since its digest_ is alignas(32) on DMA-SHA variants; md5_set_ stays last so buf_ packs tightly.
   md5::MD5Digest md5_{};
   esp_ota_handle_t update_handle_{0};
-  const esp_partition_t *partition_;
+  const esp_partition_t *partition_{nullptr};
+  size_t written_{0};  // Bytes handed to esp_ota_write()
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
+  size_t erased_end_{0};  // Erased up to this partition offset; must stay >= written_
+#endif
   char expected_bin_md5_[32];
   bool md5_set_{false};
 #ifdef USE_OTA_PARTITIONS

esphome/components/ota/ota_bootloader_esp_idf.cpp

@@ -1,6 +1,7 @@
 #ifdef USE_ESP32
 #include "ota_backend_esp_idf.h"
 
+#include "esphome/components/watchdog/watchdog.h"
 #include "esphome/core/defines.h"
 
 #ifdef USE_OTA_PARTITIONS
@@ -69,12 +70,20 @@ OTAResponseTypes IDFOTABackend::setup_bootloader_staging_() {
     return OTA_RESPONSE_ERROR_BOOTLOADER_VERIFY;
   }
   // Erase full size of the bootloader partition in the staging partition
-  // to avoid copying old data to the bootloader partition later
+  // to avoid copying old data to the bootloader partition later. Up to
+  // ESP_BOOTLOADER_SIZE of blocking erase; widen the WDT for its duration.
+  watchdog::WatchdogManager watchdog(15000);
   esp_err_t err = esp_partition_erase_range(this->partition_, 0, this->bootloader_part_->size);
   if (err != ESP_OK) {
     ESP_LOGW(TAG, "esp_partition_erase_range failed (err=0x%X)", err);
     // No critical error, don't return
   }
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
+  if (err == ESP_OK) {
+    // Skip re-erasing the pre-erased staging region in erase_ahead_()
+    this->erased_end_ = this->bootloader_part_->size;
+  }
+#endif
   err = esp_ota_set_final_partition(this->update_handle_, this->bootloader_part_, false);
   if (err != ESP_OK) {
     esp_ota_abort(this->update_handle_);

esphome/components/ota/ota_signature_esp_idf.cpp

@@ -211,7 +211,7 @@ bool rsa_pss_verify(uint8_t *block, const uint8_t *digest) {
 bool IDFOTABackend::verify_signed_image_(const esp_partition_t *incoming) {
   // Verification re-hashes the full image (after esp_ota_end already did one
   // pass), which can approach the task WDT budget on a large app. Extend it for
-  // the duration, mirroring the erase budget in begin().
+  // the duration, scaled to the image size over a 15 s floor.
   const uint32_t verify_budget_ms = 15000 + (incoming->size >> 10) * 10;
   watchdog::WatchdogManager watchdog(verify_budget_ms);
 

tests/components/ota/test_erase_ahead.cpp

@@ -0,0 +1,41 @@
+// Pins the lazy erase-ahead arithmetic used by the ESP-IDF OTA backend: the
+// erased watermark must always cover the write end, stay 64 KiB block-aligned
+// until the clamp, and never exceed the partition.
+
+#include <gtest/gtest.h>
+
+#include "esphome/components/ota/ota_backend.h"
+
+namespace esphome::ota::testing {
+
+static constexpr size_t BLOCK = 64 * 1024;
+static constexpr size_t PART = 1835008;  // 0x1C0000, a real app slot size
+
+TEST(NextEraseEnd, FirstWriteRoundsUpToOneBlock) { EXPECT_EQ(next_erase_end(1024, PART), BLOCK); }
+
+TEST(NextEraseEnd, ExactBlockBoundaryDoesNotOverErase) { EXPECT_EQ(next_erase_end(BLOCK, PART), BLOCK); }
+
+TEST(NextEraseEnd, StraddlingWriteCoversNextBlock) { EXPECT_EQ(next_erase_end(BLOCK + 1, PART), 2 * BLOCK); }
+
+TEST(NextEraseEnd, ClampsToPartitionEnd) {
+  // Partition sizes are sector multiples but not always block multiples
+  constexpr size_t part = 27 * BLOCK + 4096;
+  EXPECT_EQ(next_erase_end(27 * BLOCK + 1, part), part);
+  EXPECT_EQ(next_erase_end(part, part), part);
+}
+
+// Bootloader staging seeds erased_end_ mid-block (e.g. 0x8000); the target for
+// a write past that seed must still cover the write end.
+TEST(NextEraseEnd, MidBlockSeedStillCovered) { EXPECT_EQ(next_erase_end(0x8000 + 1024, PART), BLOCK); }
+
+TEST(NextEraseEnd, SweepAlwaysCoversWriteEndWithinPartition) {
+  for (size_t end = 1; end <= PART; end += 4093) {
+    const size_t erased = next_erase_end(end, PART);
+    ASSERT_GE(erased, end);
+    ASSERT_LE(erased, PART);
+    // Block-aligned unless clamped at the partition end
+    ASSERT_TRUE(erased == PART || erased % BLOCK == 0);
+  }
+}
+
+}  // namespace esphome::ota::testing