From 54106179a11ec9ff4586f73f8246507165ba9390 Mon Sep 17 00:00:00 2001
From: Oxan van Leeuwen <oxan@oxanvanleeuwen.nl>
Date: Wed, 1 Dec 2021 21:05:42 +0100
Subject: [PATCH] Set ESP32 watchdog to loop task (#2846)

---
 esphome/components/esp32/__init__.py |  7 ++++++
 esphome/components/esp32/core.cpp    | 35 ++++++++++++++++------------
 esphome/components/esp8266/core.cpp  |  1 +
 esphome/core/application.h           |  2 ++
 esphome/core/hal.h                   |  1 +
 sdkconfig.defaults                   |  4 ++++
 6 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/esphome/components/esp32/__init__.py b/esphome/components/esp32/__init__.py
index 1c249476e7..d6f1180aa7 100644
--- a/esphome/components/esp32/__init__.py
+++ b/esphome/components/esp32/__init__.py
@@ -311,9 +311,16 @@ async def to_code(config):
         )
         add_idf_sdkconfig_option("CONFIG_COMPILER_OPTIMIZATION_DEFAULT", False)
         add_idf_sdkconfig_option("CONFIG_COMPILER_OPTIMIZATION_SIZE", True)
+
         # Increase freertos tick speed from 100Hz to 1kHz so that delay() resolution is 1ms
         add_idf_sdkconfig_option("CONFIG_FREERTOS_HZ", 1000)
 
+        # Setup watchdog
+        add_idf_sdkconfig_option("CONFIG_ESP_TASK_WDT", True)
+        add_idf_sdkconfig_option("CONFIG_ESP_TASK_WDT_PANIC", True)
+        add_idf_sdkconfig_option("CONFIG_ESP_TASK_WDT_CHECK_IDLE_TASK_CPU0", False)
+        add_idf_sdkconfig_option("CONFIG_ESP_TASK_WDT_CHECK_IDLE_TASK_CPU1", False)
+
         cg.add_platformio_option("board_build.partitions", "partitions.csv")
 
         for name, value in conf[CONF_SDKCONFIG_OPTIONS].items():
diff --git a/esphome/components/esp32/core.cpp b/esphome/components/esp32/core.cpp
index 359999120f..a9756b41cd 100644
--- a/esphome/components/esp32/core.cpp
+++ b/esphome/components/esp32/core.cpp
@@ -6,12 +6,17 @@
 #include <freertos/FreeRTOS.h>
 #include <freertos/task.h>
 #include <esp_idf_version.h>
+#include <esp_task_wdt.h>
 #include <soc/rtc.h>
 
 #if ESP_IDF_VERSION_MAJOR >= 4
 #include <hal/cpu_hal.h>
 #endif
 
+#ifdef USE_ARDUINO
+#include <esp32-hal.h>
+#endif
+
 void setup();
 void loop();
 
@@ -29,24 +34,24 @@ void arch_restart() {
     yield();
   }
 }
-void IRAM_ATTR HOT arch_feed_wdt() {
-#ifdef USE_ARDUINO
-#if CONFIG_ARDUINO_RUNNING_CORE == 0
-#ifdef CONFIG_TASK_WDT_CHECK_IDLE_TASK_CPU0
-  // ESP32 uses "Task Watchdog" which is hooked to the FreeRTOS idle task.
-  // To cause the Watchdog to be triggered we need to put the current task
-  // to sleep to get the idle task scheduled.
-  delay(1);
-#endif
-#endif
-#endif  // USE_ARDUINO
 
-#ifdef USE_ESP_IDF
-#ifdef CONFIG_TASK_WDT_CHECK_IDLE_TASK_CPU0
-  delay(1);
+void arch_init() {
+  // Enable the task watchdog only on the loop task (from which we're currently running)
+#if defined(USE_ESP_IDF)
+  esp_task_wdt_add(nullptr);
+  // Idle task watchdog is disabled on ESP-IDF
+#elif defined(USE_ARDUINO)
+  enableLoopWDT();
+  // Disable idle task watchdog on the core we're using (Arduino pins the process to a core)
+#if CONFIG_ARDUINO_RUNNING_CORE == 0
+  disableCore0WDT();
+#endif
+#if CONFIG_ARDUINO_RUNNING_CORE == 1
+  disableCore1WDT();
+#endif
 #endif
-#endif  // USE_ESP_IDF
 }
+void IRAM_ATTR HOT arch_feed_wdt() { esp_task_wdt_reset(); }
 
 uint8_t progmem_read_byte(const uint8_t *addr) { return *addr; }
 uint32_t arch_get_cpu_cycle_count() {
diff --git a/esphome/components/esp8266/core.cpp b/esphome/components/esp8266/core.cpp
index 137d4382b4..828d71a3bd 100644
--- a/esphome/components/esp8266/core.cpp
+++ b/esphome/components/esp8266/core.cpp
@@ -20,6 +20,7 @@ void arch_restart() {
     yield();
   }
 }
+void arch_init() {}
 void IRAM_ATTR HOT arch_feed_wdt() {
   ESP.wdtFeed();  // NOLINT(readability-static-accessed-through-instance)
 }
diff --git a/esphome/core/application.h b/esphome/core/application.h
index f4fe571490..2a20793c19 100644
--- a/esphome/core/application.h
+++ b/esphome/core/application.h
@@ -5,6 +5,7 @@
 #include "esphome/core/defines.h"
 #include "esphome/core/preferences.h"
 #include "esphome/core/component.h"
+#include "esphome/core/hal.h"
 #include "esphome/core/helpers.h"
 #include "esphome/core/scheduler.h"
 
@@ -47,6 +48,7 @@ namespace esphome {
 class Application {
  public:
   void pre_setup(const std::string &name, const char *compilation_time, bool name_add_mac_suffix) {
+    arch_init();
     this->name_add_mac_suffix_ = name_add_mac_suffix;
     if (name_add_mac_suffix) {
       this->name_ = name + "-" + get_mac_address().substr(6);
diff --git a/esphome/core/hal.h b/esphome/core/hal.h
index a86dbf2534..034f9d692f 100644
--- a/esphome/core/hal.h
+++ b/esphome/core/hal.h
@@ -39,6 +39,7 @@ uint32_t micros();
 void delay(uint32_t ms);
 void delayMicroseconds(uint32_t us);  // NOLINT(readability-identifier-naming)
 void __attribute__((noreturn)) arch_restart();
+void arch_init();
 void arch_feed_wdt();
 uint32_t arch_get_cpu_cycle_count();
 uint32_t arch_get_cpu_freq_hz();
diff --git a/sdkconfig.defaults b/sdkconfig.defaults
index 26db4705b8..72ca3f6e9c 100644
--- a/sdkconfig.defaults
+++ b/sdkconfig.defaults
@@ -9,6 +9,10 @@ CONFIG_PARTITION_TABLE_CUSTOM=y
 #CONFIG_PARTITION_TABLE_CUSTOM_FILENAME="partitions.csv"
 CONFIG_PARTITION_TABLE_SINGLE_APP=n
 CONFIG_FREERTOS_HZ=1000
+CONFIG_ESP_TASK_WDT=y
+CONFIG_ESP_TASK_WDT_PANIC=y
+CONFIG_ESP_TASK_WDT_CHECK_IDLE_TASK_CPU0=n
+CONFIG_ESP_TASK_WDT_CHECK_IDLE_TASK_CPU1=n
 
 # esp32_ble
 CONFIG_BT_ENABLED=y