io_stress/
stress.rs

1// Copyright 2026 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use anyhow::{Context as _, Result};
6pub use fidl_fuchsia_memorypressure::Level as MemoryLevel;
7use fidl_fuchsia_memorypressure::{ProviderMarker, WatcherMarker, WatcherRequest};
8use fuchsia_component::client::connect_to_protocol;
9use futures::StreamExt as _;
10use std::sync::Arc;
11use std::sync::atomic::{AtomicBool, Ordering};
12use std::thread;
13
14pub struct CpuStressor {
15    stop_signal: Arc<AtomicBool>,
16    threads: Vec<thread::JoinHandle<()>>,
17}
18
19impl CpuStressor {
20    pub fn start(num_cores: usize) -> Self {
21        log::info!("Spawning {} background CPU burn threads...", num_cores);
22        let stop_signal = Arc::new(AtomicBool::new(false));
23        let mut threads = Vec::new();
24
25        for i in 0..num_cores {
26            let stop = stop_signal.clone();
27            threads.push(thread::spawn(move || {
28                log::debug!("CPU burn thread {} started.", i);
29                while !stop.load(Ordering::Relaxed) {
30                    let mut _x = 0;
31                    for _ in 0..1000 {
32                        _x = std::hint::black_box(_x + 1);
33                    }
34                }
35                log::debug!("CPU burn thread {} stopped.", i);
36            }));
37        }
38
39        Self { stop_signal, threads }
40    }
41}
42
43impl Drop for CpuStressor {
44    fn drop(&mut self) {
45        log::info!("Stopping background CPU burn threads...");
46        self.stop_signal.store(true, Ordering::Relaxed);
47        for handle in self.threads.drain(..) {
48            let _ = handle.join();
49        }
50        log::info!("All CPU burn threads stopped.");
51    }
52}
53
54pub fn set_memory_pressure(level: fidl_fuchsia_memorypressure::Level) -> Result<()> {
55    log::info!("Signaling memory pressure level to target: {:?}", level);
56    let debug_pressure = match connect_to_protocol::<fidl_fuchsia_memory_debug::MemoryPressureMarker>(
57    ) {
58        Ok(p) => p,
59        Err(e) => {
60            log::warn!(
61                "Could not connect to fuchsia.memory.debug.MemoryPressure (routing might be missing): {:?}",
62                e
63            );
64            return Ok(());
65        }
66    };
67    if let Err(e) = debug_pressure.signal(level) {
68        log::warn!("Failed to signal memory pressure (PEER_CLOSED?): {:?}", e);
69    }
70    Ok(())
71}
72
73pub struct MemoryPressureGuard;
74
75impl MemoryPressureGuard {
76    pub fn new(level: fidl_fuchsia_memorypressure::Level) -> Result<Self> {
77        set_memory_pressure(level)?;
78        Ok(Self)
79    }
80}
81
82impl Drop for MemoryPressureGuard {
83    fn drop(&mut self) {
84        if let Err(e) = set_memory_pressure(fidl_fuchsia_memorypressure::Level::Normal) {
85            log::error!("Failed to reset memory pressure to Normal on drop: {:?}", e);
86        }
87    }
88}
89
90pub async fn get_cpu_load() -> Result<f32> {
91    let stats = connect_to_protocol::<fidl_fuchsia_kernel::StatsMarker>()
92        .context("Failed to connect to fuchsia.kernel.Stats")?;
93    let duration = zx::MonotonicDuration::from_millis(200).into_nanos();
94    let load_vec = stats.get_cpu_load(duration).await.context("FIDL call to GetCpuLoad failed")?;
95    if load_vec.is_empty() {
96        return Ok(0.0);
97    }
98    let total_load: f32 = load_vec.iter().sum();
99    Ok(total_load / load_vec.len() as f32)
100}
101
102pub async fn wait_for_cooldown(
103    stats: &MemoryPressureStats,
104    max_wait: std::time::Duration,
105) -> Result<()> {
106    log::info!("Starting dynamic cooldown... max wait: {:?}", max_wait);
107    let start = std::time::Instant::now();
108    let check_interval = std::time::Duration::from_secs(1);
109    let mut last_cpu_load = 0.0;
110
111    while start.elapsed() < max_wait {
112        let mem_level = stats.current_level();
113        let cpu_load = match get_cpu_load().await {
114            Ok(load) => load,
115            Err(e) => {
116                log::warn!("Failed to query CPU load: {:?}", e);
117                0.0
118            }
119        };
120        last_cpu_load = cpu_load;
121
122        log::info!(
123            "Cooldown check: mem_level = {:?}, cpu_load = {:.2}% (elapsed: {:?})",
124            mem_level,
125            cpu_load,
126            start.elapsed()
127        );
128
129        if mem_level == fidl_fuchsia_memorypressure::Level::Normal && cpu_load < 40.0 {
130            log::info!("System settled after {:?}", start.elapsed());
131            return Ok(());
132        }
133
134        fuchsia_async::Timer::new(check_interval).await;
135    }
136
137    anyhow::bail!(
138        "System failed to settle within {:?}. Final state: mem_level={:?}, cpu_load={:.2}%",
139        max_wait,
140        stats.current_level(),
141        last_cpu_load
142    );
143}
144
145use std::sync::Mutex;
146
147pub struct MemoryPressureStats {
148    inner: Mutex<MemoryPressureStatsInner>,
149}
150
151struct MemoryPressureStatsInner {
152    max_level: MemoryLevel,
153    current_level: MemoryLevel,
154}
155
156impl MemoryPressureStats {
157    pub fn new() -> Self {
158        Self {
159            inner: Mutex::new(MemoryPressureStatsInner {
160                max_level: MemoryLevel::Normal,
161                current_level: MemoryLevel::Normal,
162            }),
163        }
164    }
165
166    pub fn update(&self, new_level: MemoryLevel) {
167        let mut inner = self.inner.lock().unwrap();
168        inner.max_level = std::cmp::max(inner.max_level, new_level);
169        inner.current_level = new_level;
170    }
171
172    pub fn max_level(&self) -> MemoryLevel {
173        self.inner.lock().unwrap().max_level
174    }
175
176    pub fn current_level(&self) -> MemoryLevel {
177        self.inner.lock().unwrap().current_level
178    }
179
180    pub fn reset(&self) {
181        let mut inner = self.inner.lock().unwrap();
182        inner.max_level = MemoryLevel::Normal;
183        inner.current_level = MemoryLevel::Normal;
184    }
185}
186
187pub async fn watch_memory_pressure(stats: Arc<MemoryPressureStats>) -> anyhow::Result<()> {
188    let provider = connect_to_protocol::<ProviderMarker>()
189        .context("Failed to connect to fuchsia.memorypressure.Provider")?;
190
191    let (watcher_client, mut watcher_requests) =
192        fidl::endpoints::create_request_stream::<WatcherMarker>();
193    provider
194        .register_watcher(watcher_client)
195        .context("Failed to register watcher with MemoryPressure Provider")?;
196
197    log::info!("Registered as a fuchsia.memorypressure/Watcher");
198
199    while let Some(request) = watcher_requests.next().await {
200        match request {
201            Ok(WatcherRequest::OnLevelChanged { level, responder }) => {
202                log::info!("Memory pressure level changed to: {:?}", level);
203                stats.update(level);
204                if let Err(e) = responder.send() {
205                    log::error!("Failed to respond to memory pressure event: {:?}", e);
206                    break;
207                }
208            }
209            Err(e) => {
210                log::error!("Error reading memory pressure watcher stream: {:?}", e);
211                break;
212            }
213        }
214    }
215    Ok(())
216}
io_stress/stress.rs

io_stress/
stress.rs
