Fix get_systimer_count racy behaviour

esp-wifi/src/tasks.rs

@@ -35,6 +35,7 @@ pub extern "C" fn worker_task2() {
     loop {
         let mut to_run = SimpleQueue::<_, 20>::new();
 
+        let current_timestamp = get_systimer_count();
         critical_section::with(|_| unsafe {
             let current_timestamp = get_systimer_count();
             memory_fence();

esp-wifi/src/timer_esp32.rs

@@ -25,10 +25,22 @@ static TIMER1: Mutex<RefCell<Option<Timer<Timer0<TIMG1>>>>> = Mutex::new(RefCell
 
 static TIMER_OVERFLOWS: AtomicU32 = AtomicU32::new(0);
 
+/// This function must not be called in a critical section. Doing so may return an incorrect value.
 pub fn get_systimer_count() -> u64 {
-    let overflow = (TIMER_OVERFLOWS.load(Ordering::Relaxed) as u64) << 32;
-    let counter = xtensa_lx::timer::get_cycle_count() as u64;
-    (overflow + counter) * 40_000_000 / 240_000_000
+    // We read the cycle counter twice to detect overflows.
+    // If we don't detect an overflow, we use the TIMER_OVERFLOWS count we read between.
+    // If we detect an overflow, we read the TIMER_OVERFLOWS count again to make sure we use the
+    // value after the overflow has been handled.
+
+    let counter_before = xtensa_lx::timer::get_cycle_count();
+    let mut overflow = TIMER_OVERFLOWS.load(Ordering::Relaxed);
+    let counter_after = xtensa_lx::timer::get_cycle_count();
+
+    if counter_after < counter_before {
+        overflow = TIMER_OVERFLOWS.load(Ordering::Relaxed);
+    }
+
+    (((overflow as u64) << 32) + counter_after as u64) * 40_000_000 / 240_000_000
 }
 
 pub fn setup_timer_isr(timg1_timer0: Timer<Timer0<TIMG1>>) {

esp-wifi/src/timer_esp32s2.rs

@@ -25,10 +25,22 @@ static TIMER1: Mutex<RefCell<Option<Timer<Timer0<TIMG1>>>>> = Mutex::new(RefCell
 
 static TIMER_OVERFLOWS: AtomicU32 = AtomicU32::new(0);
 
+/// This function must not be called in a critical section. Doing so may return an incorrect value.
 pub fn get_systimer_count() -> u64 {
-    let overflow = (TIMER_OVERFLOWS.load(Ordering::Relaxed) as u64) << 32;
-    let counter = xtensa_lx::timer::get_cycle_count() as u64;
-    (overflow + counter) * 40_000_000 / 240_000_000
+    // We read the cycle counter twice to detect overflows.
+    // If we don't detect an overflow, we use the TIMER_OVERFLOWS count we read between.
+    // If we detect an overflow, we read the TIMER_OVERFLOWS count again to make sure we use the
+    // value after the overflow has been handled.
+
+    let counter_before = xtensa_lx::timer::get_cycle_count();
+    let mut overflow = TIMER_OVERFLOWS.load(Ordering::Relaxed);
+    let counter_after = xtensa_lx::timer::get_cycle_count();
+
+    if counter_after < counter_before {
+        overflow = TIMER_OVERFLOWS.load(Ordering::Relaxed);
+    }
+
+    (((overflow as u64) << 32) + counter_after as u64) * 40_000_000 / 240_000_000
 }
 
 pub fn setup_timer_isr(timg1_timer0: Timer<Timer0<TIMG1>>) {

esp-wifi/src/timer_esp32s3.rs

@@ -8,6 +8,7 @@ use esp32s3_hal::{
     peripherals::{self, TIMG1},
     prelude::*,
     timer::{Timer, Timer0},
+    xtensa_lx,
 };
 
 use crate::preempt::preempt::task_switch;
@@ -23,10 +24,22 @@ static TIMER1: Mutex<RefCell<Option<Timer<Timer0<TIMG1>>>>> = Mutex::new(RefCell
 
 static TIMER_OVERFLOWS: AtomicU32 = AtomicU32::new(0);
 
+/// This function must not be called in a critical section. Doing so may return an incorrect value.
 pub fn get_systimer_count() -> u64 {
-    let overflow = (TIMER_OVERFLOWS.load(Ordering::Relaxed) as u64) << 32;
-    let counter = esp32s3_hal::xtensa_lx::timer::get_cycle_count() as u64;
-    (overflow + counter) * 40_000_000 / 240_000_000
+    // We read the cycle counter twice to detect overflows.
+    // If we don't detect an overflow, we use the TIMER_OVERFLOWS count we read between.
+    // If we detect an overflow, we read the TIMER_OVERFLOWS count again to make sure we use the
+    // value after the overflow has been handled.
+
+    let counter_before = xtensa_lx::timer::get_cycle_count();
+    let mut overflow = TIMER_OVERFLOWS.load(Ordering::Relaxed);
+    let counter_after = xtensa_lx::timer::get_cycle_count();
+
+    if counter_after < counter_before {
+        overflow = TIMER_OVERFLOWS.load(Ordering::Relaxed);
+    }
+
+    (((overflow as u64) << 32) + counter_after as u64) * 40_000_000 / 240_000_000
 }
 
 pub fn setup_timer_isr(timg1_timer0: Timer<Timer0<TIMG1>>) {