Avoid AtomicU64 if possible, fix race condition (esp-rs#278)

bugadani · bjoernQ · commit 3adfbea0f8da · 2024-05-24T13:52:10.000+02:00
* Xtensa: avoid AtomicU64 in get_systimer_count

* Fix get_systimer_count racy behaviour

* Improve code consistency
diff --git a/esp-wifi/src/tasks.rs b/esp-wifi/src/tasks.rs
@@ -35,8 +35,8 @@ 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();
             for i in 0..TIMERS.len() {
                 if let Some(ref mut timer) = TIMERS[i] {
diff --git a/esp-wifi/src/timer_esp32.rs b/esp-wifi/src/timer_esp32.rs
@@ -1,15 +1,14 @@
 use core::cell::RefCell;
 
-use atomic_polyfill::{AtomicU64, Ordering};
+use atomic_polyfill::{AtomicU32, Ordering};
 use critical_section::Mutex;
-use esp32_hal::xtensa_lx;
-use esp32_hal::xtensa_lx_rt;
-use esp32_hal::xtensa_lx_rt::exception::Context;
+use esp32_hal::trapframe::TrapFrame;
 use esp32_hal::{
     interrupt,
     peripherals::{self, TIMG1},
     prelude::*,
     timer::{Timer, Timer0},
+    xtensa_lx, xtensa_lx_rt,
 };
 
 use crate::preempt::preempt::task_switch;
@@ -23,16 +22,24 @@ const TIMER_DELAY: fugit::HertzU32 = fugit::HertzU32::from_raw(crate::CONFIG.tic
 
 static TIMER1: Mutex<RefCell<Option<Timer<Timer0<TIMG1>>>>> = Mutex::new(RefCell::new(None));
 
-static TIME: AtomicU64 = AtomicU64::new(0);
+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 {
-    TIME.load(Ordering::Relaxed) + read_timer_value()
-}
+    // 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);
+    }
 
-#[inline(always)]
-fn read_timer_value() -> u64 {
-    let value = xtensa_lx::timer::get_cycle_count() as u64;
-    value * 40_000_000 / 240_000_000
+    (((overflow as u64) << 32) + counter_after as u64) * 40_000_000 / 240_000_000
 }
 
 pub fn setup_timer_isr(timg1_timer0: Timer<Timer0<TIMG1>>) {
@@ -74,7 +81,7 @@ pub fn setup_timer_isr(timg1_timer0: Timer<Timer0<TIMG1>>) {
     xtensa_lx::timer::set_ccompare0(0xffffffff);
 
     unsafe {
-        let enabled = esp32_hal::xtensa_lx::interrupt::disable();
+        let enabled = xtensa_lx::interrupt::disable();
         xtensa_lx::interrupt::enable_mask(
             1 << 6 // Timer0
             | 1 << 29 // Software1
@@ -104,7 +111,7 @@ fn Software0(_level: u32) {
 #[allow(non_snake_case)]
 #[no_mangle]
 fn Timer0(_level: u32) {
-    TIME.fetch_add(0x1_0000_0000 * 40_000_000 / 240_000_000, Ordering::Relaxed);
+    TIMER_OVERFLOWS.fetch_add(1, Ordering::Relaxed);
 
     xtensa_lx::timer::set_ccompare0(0xffffffff);
 }
@@ -166,7 +173,7 @@ fn BT_BB() {
 }
 
 #[interrupt]
-fn TG1_T0_LEVEL(context: &mut Context) {
+fn TG1_T0_LEVEL(context: &mut TrapFrame) {
     task_switch(context);
 
     critical_section::with(|cs| {
@@ -181,7 +188,7 @@ fn TG1_T0_LEVEL(context: &mut Context) {
 
 #[allow(non_snake_case)]
 #[no_mangle]
-fn Software1(_level: u32, context: &mut Context) {
+fn Software1(_level: u32, context: &mut TrapFrame) {
     let intr = 1 << 29;
     unsafe {
         core::arch::asm!("wsr.intclear  {0}", in(reg) intr, options(nostack));
diff --git a/esp-wifi/src/timer_esp32s2.rs b/esp-wifi/src/timer_esp32s2.rs
@@ -1,15 +1,14 @@
 use core::cell::RefCell;
 
-use atomic_polyfill::{AtomicU64, Ordering};
+use atomic_polyfill::{AtomicU32, Ordering};
 use critical_section::Mutex;
-use esp32s2_hal::xtensa_lx;
-use esp32s2_hal::xtensa_lx_rt;
-use esp32s2_hal::xtensa_lx_rt::exception::Context;
+use esp32s2_hal::trapframe::TrapFrame;
 use esp32s2_hal::{
     interrupt,
     peripherals::{self, TIMG1},
     prelude::*,
     timer::{Timer, Timer0},
+    xtensa_lx, xtensa_lx_rt,
 };
 
 use crate::preempt::preempt::task_switch;
@@ -23,16 +22,24 @@ const TIMER_DELAY: fugit::HertzU32 = fugit::HertzU32::from_raw(crate::CONFIG.tic
 
 static TIMER1: Mutex<RefCell<Option<Timer<Timer0<TIMG1>>>>> = Mutex::new(RefCell::new(None));
 
-static TIME: AtomicU64 = AtomicU64::new(0);
+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 {
-    TIME.load(Ordering::Relaxed) + read_timer_value()
-}
+    // 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);
+    }
 
-#[inline(always)]
-fn read_timer_value() -> u64 {
-    let value = xtensa_lx::timer::get_cycle_count() as u64;
-    value * 40_000_000 / 240_000_000
+    (((overflow as u64) << 32) + counter_after as u64) * 40_000_000 / 240_000_000
 }
 
 pub fn setup_timer_isr(timg1_timer0: Timer<Timer0<TIMG1>>) {
@@ -63,7 +70,7 @@ pub fn setup_timer_isr(timg1_timer0: Timer<Timer0<TIMG1>>) {
     xtensa_lx::timer::set_ccompare0(0xffffffff);
 
     unsafe {
-        let enabled = esp32s2_hal::xtensa_lx::interrupt::disable();
+        let enabled = xtensa_lx::interrupt::disable();
         xtensa_lx::interrupt::enable_mask(
             1 << 6 // Timer0
             | 1 << 29 // Software1
@@ -78,7 +85,7 @@ pub fn setup_timer_isr(timg1_timer0: Timer<Timer0<TIMG1>>) {
 #[allow(non_snake_case)]
 #[no_mangle]
 fn Timer0(_level: u32) {
-    TIME.fetch_add(0x1_0000_0000 * 40_000_000 / 240_000_000, Ordering::Relaxed);
+    TIMER_OVERFLOWS.fetch_add(1, Ordering::Relaxed);
 
     xtensa_lx::timer::set_ccompare0(0xffffffff);
 }
@@ -115,7 +122,7 @@ fn WIFI_PWR() {
 }
 
 #[interrupt]
-fn TG1_T0_LEVEL(context: &mut Context) {
+fn TG1_T0_LEVEL(context: &mut TrapFrame) {
     task_switch(context);
 
     critical_section::with(|cs| {
@@ -130,7 +137,7 @@ fn TG1_T0_LEVEL(context: &mut Context) {
 
 #[allow(non_snake_case)]
 #[no_mangle]
-fn Software1(_level: u32, context: &mut Context) {
+fn Software1(_level: u32, context: &mut TrapFrame) {
     let intr = 1 << 29;
     unsafe {
         core::arch::asm!("wsr.intclear  {0}", in(reg) intr, options(nostack));
diff --git a/esp-wifi/src/timer_esp32s3.rs b/esp-wifi/src/timer_esp32s3.rs
@@ -1,13 +1,14 @@
 use core::cell::RefCell;
 
-use atomic_polyfill::{AtomicU64, Ordering};
+use atomic_polyfill::{AtomicU32, Ordering};
 use critical_section::Mutex;
 use esp32s3_hal::trapframe::TrapFrame;
 use esp32s3_hal::{
     interrupt,
     peripherals::{self, TIMG1},
     prelude::*,
     timer::{Timer, Timer0},
+    xtensa_lx, xtensa_lx_rt,
 };
 
 use crate::preempt::preempt::task_switch;
@@ -21,16 +22,24 @@ const TIMER_DELAY: fugit::HertzU32 = fugit::HertzU32::from_raw(crate::CONFIG.tic
 
 static TIMER1: Mutex<RefCell<Option<Timer<Timer0<TIMG1>>>>> = Mutex::new(RefCell::new(None));
 
-static TIME: AtomicU64 = AtomicU64::new(0);
+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 {
-    TIME.load(Ordering::Relaxed) + read_timer_value()
-}
+    // 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);
+    }
 
-#[inline(always)]
-fn read_timer_value() -> u64 {
-    let value = esp32s3_hal::xtensa_lx::timer::get_cycle_count() as u64;
-    value * 40_000_000 / 240_000_000
+    (((overflow as u64) << 32) + counter_after as u64) * 40_000_000 / 240_000_000
 }
 
 pub fn setup_timer_isr(timg1_timer0: Timer<Timer0<TIMG1>>) {
@@ -73,12 +82,12 @@ pub fn setup_timer_isr(timg1_timer0: Timer<Timer0<TIMG1>>) {
     esp32s3_hal::xtensa_lx::timer::set_ccompare0(0xffffffff);
 
     unsafe {
-        let enabled = esp32s3_hal::xtensa_lx::interrupt::disable();
-        esp32s3_hal::xtensa_lx::interrupt::enable_mask(
+        let enabled = xtensa_lx::interrupt::disable();
+        xtensa_lx::interrupt::enable_mask(
             1 << 6 // Timer0
             | 1 << 29 // Software1
-                | esp32s3_hal::xtensa_lx_rt::interrupt::CpuInterruptLevel::Level2.mask()
-                | esp32s3_hal::xtensa_lx_rt::interrupt::CpuInterruptLevel::Level6.mask() | enabled,
+                | xtensa_lx_rt::interrupt::CpuInterruptLevel::Level2.mask()
+                | xtensa_lx_rt::interrupt::CpuInterruptLevel::Level6.mask() | enabled,
         );
     }
 
@@ -88,7 +97,7 @@ pub fn setup_timer_isr(timg1_timer0: Timer<Timer0<TIMG1>>) {
 #[allow(non_snake_case)]
 #[no_mangle]
 fn Timer0(_level: u32) {
-    TIME.fetch_add(0x1_0000_0000 * 40_000_000 / 240_000_000, Ordering::Relaxed);
+    TIMER_OVERFLOWS.fetch_add(1, Ordering::Relaxed);
 
     esp32s3_hal::xtensa_lx::timer::set_ccompare0(0xffffffff);
 }
