From ba7da69e2597c2b49eb43e8bcd609d66c1a28e1c Mon Sep 17 00:00:00 2001 From: Kang Date: Wed, 7 Oct 2026 23:22:40 +0800 Subject: [PATCH] refactor(session): replace fetch_update with compare-exchange loops for atomic operations - Simplified atomic operations in SessionCommandQueueMetrics by replacing fetch_update with custom compare-exchange loops for better clarity and performance. - Added new utility functions for saturating atomic addition and subtraction for both usize and u64 types. - Introduced tests to ensure that counters saturate correctly at their limits in both session command metrics and terminal presentation step counters. --- src-tauri/src/core/mcp/host.rs | 30 ++++++++-- src-tauri/src/core/session.rs | 101 +++++++++++++++++++++++---------- 2 files changed, 96 insertions(+), 35 deletions(-) diff --git a/src-tauri/src/core/mcp/host.rs b/src-tauri/src/core/mcp/host.rs index 80ba679a..39f65b02 100644 --- a/src-tauri/src/core/mcp/host.rs +++ b/src-tauri/src/core/mcp/host.rs @@ -159,12 +159,18 @@ impl McpTerminalPresentation { } fn next_spec(&self) -> McpTerminalPresentationSpec { - let step_index = self - .next_step_index - .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |current| { - Some(current.saturating_add(1)) - }) - .unwrap_or_else(|current| current); + let mut step_index = self.next_step_index.load(Ordering::SeqCst); + loop { + match self.next_step_index.compare_exchange_weak( + step_index, + step_index.saturating_add(1), + Ordering::SeqCst, + Ordering::SeqCst, + ) { + Ok(_) => break, + Err(actual) => step_index = actual, + } + } McpTerminalPresentationSpec { step_index, max_lines: self.max_lines, @@ -2134,4 +2140,16 @@ mod tests { assert_eq!(second.next_spec().step_index, 0); assert_eq!(second.next_spec().max_lines, 20); } + + #[test] + fn terminal_presentation_step_counter_saturates() { + let presentation = McpTerminalPresentation::new(10); + presentation + .next_step_index + .store(u16::MAX - 1, Ordering::SeqCst); + + assert_eq!(presentation.next_spec().step_index, u16::MAX - 1); + assert_eq!(presentation.next_spec().step_index, u16::MAX); + assert_eq!(presentation.next_spec().step_index, u16::MAX); + } } diff --git a/src-tauri/src/core/session.rs b/src-tauri/src/core/session.rs index 22bcf5a2..e1960047 100644 --- a/src-tauri/src/core/session.rs +++ b/src-tauri/src/core/session.rs @@ -317,18 +317,8 @@ impl SessionCommandQueueMetrics { } fn reserve_enqueue(&self) -> usize { - self.queued_commands - .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |value| { - Some(value.saturating_add(1)) - }) - .unwrap_or(u64::MAX); - let previous = self - .current_pending - .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |value| { - Some(value.saturating_add(1)) - }) - .unwrap_or(usize::MAX); - previous.saturating_add(1) + saturating_atomic_add_u64(&self.queued_commands, 1); + saturating_atomic_add_usize(&self.current_pending, 1) } fn commit_enqueue(&self, pending: usize) { @@ -354,11 +344,7 @@ impl SessionCommandQueueMetrics { } fn mark_processed(&self) { - self.processed_commands - .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |value| { - Some(value.saturating_add(1)) - }) - .ok(); + saturating_atomic_add_u64(&self.processed_commands, 1); let pending = saturating_atomic_sub_usize(&self.current_pending, 1); self.relax_pressure_tier(pending); } @@ -427,22 +413,49 @@ impl SessionCommandQueueMetrics { } } +// Use compare-exchange loops to retain Rust 1.94 compatibility: try_update is newer. +fn saturating_atomic_add_usize(value: &AtomicUsize, amount: usize) -> usize { + let mut current = value.load(Ordering::Relaxed); + loop { + let next = current.saturating_add(amount); + match value.compare_exchange_weak(current, next, Ordering::Relaxed, Ordering::Relaxed) { + Ok(_) => return next, + Err(actual) => current = actual, + } + } +} + +fn saturating_atomic_add_u64(value: &AtomicU64, amount: u64) -> u64 { + let mut current = value.load(Ordering::Relaxed); + loop { + let next = current.saturating_add(amount); + match value.compare_exchange_weak(current, next, Ordering::Relaxed, Ordering::Relaxed) { + Ok(_) => return next, + Err(actual) => current = actual, + } + } +} + fn saturating_atomic_sub_usize(value: &AtomicUsize, amount: usize) -> usize { - value - .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { - Some(current.saturating_sub(amount)) - }) - .map(|previous| previous.saturating_sub(amount)) - .unwrap_or(0) + let mut current = value.load(Ordering::Relaxed); + loop { + let next = current.saturating_sub(amount); + match value.compare_exchange_weak(current, next, Ordering::Relaxed, Ordering::Relaxed) { + Ok(_) => return next, + Err(actual) => current = actual, + } + } } fn saturating_atomic_sub_u64(value: &AtomicU64, amount: u64) -> u64 { - value - .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { - Some(current.saturating_sub(amount)) - }) - .map(|previous| previous.saturating_sub(amount)) - .unwrap_or(0) + let mut current = value.load(Ordering::Relaxed); + loop { + let next = current.saturating_sub(amount); + match value.compare_exchange_weak(current, next, Ordering::Relaxed, Ordering::Relaxed) { + Ok(_) => return next, + Err(actual) => current = actual, + } + } } pub struct SessionCommandSender { @@ -1824,6 +1837,36 @@ mod tests { ); } + #[test] + fn command_queue_metrics_saturate_at_counter_limits() { + let metrics = super::SessionCommandQueueMetrics::new("queue-limits".to_string()); + metrics + .queued_commands + .store(u64::MAX - 1, Ordering::Relaxed); + metrics + .processed_commands + .store(u64::MAX - 1, Ordering::Relaxed); + metrics + .current_pending + .store(usize::MAX - 1, Ordering::Relaxed); + + assert_eq!(metrics.reserve_enqueue(), usize::MAX); + assert_eq!(metrics.reserve_enqueue(), usize::MAX); + assert_eq!(metrics.snapshot().queued_commands, u64::MAX); + metrics.mark_processed(); + metrics.mark_processed(); + assert_eq!(metrics.snapshot().processed_commands, u64::MAX); + assert_eq!(metrics.snapshot().current_pending, usize::MAX - 2); + + metrics.queued_commands.store(1, Ordering::Relaxed); + metrics.current_pending.store(1, Ordering::Relaxed); + metrics.rollback_enqueue(); + metrics.rollback_enqueue(); + metrics.mark_processed(); + assert_eq!(metrics.snapshot().queued_commands, 0); + assert_eq!(metrics.snapshot().current_pending, 0); + } + #[test] fn command_queue_metrics_roll_back_failed_send_and_clear_on_receiver_drop() { let (closed_tx, closed_rx) = session_command_channel("closed-queue");