AArch64: Add stall diagnostics for TLB shootdown and context-switch lock

Use the StallWatch helper added earlier to report a spin loop that has
run far longer than any legitimate SMP handshake should take -- the
TLB shootdown wait and the context-switch-lock acquisition -- with one
diagnostic naming which CPU and, for the shootdown case, how many
acknowledgements are outstanding. Purely diagnostic: the loop keeps
spinning exactly as before, this only makes a stuck wait visible
instead of looking like a silent hang.

Signed-off-by: Luiz Fernando Becher de Araujo <luiz.becher.araujo@gmail.com>

src/context/memory.rs

@@ -3003,8 +3003,20 @@ impl<'a, 'addrsp> Flusher<'a, 'addrsp> {
             rmm::PageFlushAll::<RmmA>::new().flush();
         }
 
+        #[cfg(target_arch = "aarch64")]
+        let mut stall_watch = crate::arch::misc::StallWatch::start(2);
+
         while self.state.ackword.load(Ordering::SeqCst) < affected_cpu_count {
             PercpuBlock::current().maybe_handle_tlb_shootdown();
+            #[cfg(target_arch = "aarch64")]
+            if stall_watch.stalled() {
+                error!(
+                    "TLB shootdown stalled on CPU {}: acknowledgements {}/{}",
+                    current_cpu_id,
+                    self.state.ackword.load(Ordering::SeqCst),
+                    affected_cpu_count
+                );
+            }
             core::hint::spin_loop();
         }
 

src/context/switch.rs

@@ -187,12 +187,19 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
     // Acquire the global lock to ensure exclusive access during context switch and avoid
     // issues that would be caused by the unsafe operations below
     // TODO: Better memory orderings?
+    #[cfg(target_arch = "aarch64")]
+    let mut lock_stall_watch = crate::arch::misc::StallWatch::start(2);
+
     while arch::CONTEXT_SWITCH_LOCK
         .compare_exchange_weak(false, true, Ordering::SeqCst, Ordering::Relaxed)
         .is_err()
     {
         hint::spin_loop();
         percpu.maybe_handle_tlb_shootdown();
+        #[cfg(target_arch = "aarch64")]
+        if lock_stall_watch.stalled() {
+            error!("context switch lock stalled on CPU {}", crate::cpu_id());
+        }
     }
 
     // Lock the previous context.