From 9fffed7160ad791e9d81d66ff7d0c0d3e085586d Mon Sep 17 00:00:00 2001
From: Skruppy <skruppy@onmars.eu>
Date: Thu, 4 Nov 2021 18:11:57 +0100
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20Prevent=20AVR=20watchdogpile=20(?=
 =?UTF-8?q?#23075)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 Marlin/Configuration_adv.h   |  3 +++
 Marlin/src/HAL/AVR/HAL.cpp   | 21 ++++++++++++++++++++-
 Marlin/src/HAL/AVR/HAL.h     |  6 +++---
 Marlin/src/inc/SanityCheck.h |  5 +++++
 4 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/Marlin/Configuration_adv.h b/Marlin/Configuration_adv.h
index ba3c2c10c9..2cd62509c5 100644
--- a/Marlin/Configuration_adv.h
+++ b/Marlin/Configuration_adv.h
@@ -4224,3 +4224,6 @@
  */
 //#define SOFT_RESET_VIA_SERIAL         // 'KILL' and '^X' commands will soft-reset the controller
 //#define SOFT_RESET_ON_KILL            // Use a digital button to soft-reset the controller after KILL
+
+// Report uncleaned reset reason from register r2 instead of MCUSR. Supported by Optiboot on AVR.
+//#define OPTIBOOT_RESET_REASON
diff --git a/Marlin/src/HAL/AVR/HAL.cpp b/Marlin/src/HAL/AVR/HAL.cpp
index 708583b262..d7bf2a6f6f 100644
--- a/Marlin/src/HAL/AVR/HAL.cpp
+++ b/Marlin/src/HAL/AVR/HAL.cpp
@@ -35,12 +35,31 @@
 // Public Variables
 // ------------------------
 
-//uint8_t MCUSR;
+// Don't initialize/override variable (which would happen in .init4)
+uint8_t reset_reason __attribute__((section(".noinit")));
 
 // ------------------------
 // Public functions
 // ------------------------
 
+__attribute__((naked))             // Don't output function pro- and epilogue
+__attribute__((used))              // Output the function, even if "not used"
+__attribute__((section(".init3"))) // Put in an early user definable section
+void HAL_save_reset_reason() {
+  #if ENABLED(OPTIBOOT_RESET_REASON)
+    __asm__ __volatile__(
+      A("STS %0, r2")
+      : "=m"(reset_reason)
+    );
+  #else
+    reset_reason = MCUSR;
+  #endif
+
+  // Clear within 16ms since WDRF bit enables a 16ms watchdog timer -> Boot loop
+  MCUSR = 0;
+  wdt_disable();
+}
+
 void HAL_init() {
   // Init Servo Pins
   #define INIT_SERVO(N) OUT_WRITE(SERVO##N##_PIN, LOW)
diff --git a/Marlin/src/HAL/AVR/HAL.h b/Marlin/src/HAL/AVR/HAL.h
index ad1f47a4e6..2217f239d6 100644
--- a/Marlin/src/HAL/AVR/HAL.h
+++ b/Marlin/src/HAL/AVR/HAL.h
@@ -91,7 +91,7 @@ typedef int8_t pin_t;
 // Public Variables
 // ------------------------
 
-//extern uint8_t MCUSR;
+extern uint8_t reset_reason;
 
 // Serial ports
 #ifdef USBCON
@@ -152,8 +152,8 @@ void HAL_init();
 
 //void _delay_ms(const int delay);
 
-inline void HAL_clear_reset_source() { MCUSR = 0; }
-inline uint8_t HAL_get_reset_source() { return MCUSR; }
+inline void HAL_clear_reset_source() { }
+inline uint8_t HAL_get_reset_source() { return reset_reason; }
 
 void HAL_reboot();
 
diff --git a/Marlin/src/inc/SanityCheck.h b/Marlin/src/inc/SanityCheck.h
index 3287262f0b..6294e3b0ba 100644
--- a/Marlin/src/inc/SanityCheck.h
+++ b/Marlin/src/inc/SanityCheck.h
@@ -2489,6 +2489,11 @@ static_assert(Y_MAX_LENGTH >= Y_BED_SIZE, "Movement bounds (Y_MIN_POS, Y_MAX_POS
   #error "An encoder button is required or SOFT_RESET_ON_KILL will reset the printer without notice!"
 #endif
 
+// Reset reason for AVR
+#if ENABLED(OPTIBOOT_RESET_REASON) && !defined(__AVR__)
+  #error "OPTIBOOT_RESET_REASON only applies to AVR."
+#endif
+
 /**
  * I2C bus
  */