starnix_core/mm/
prefetch_pool.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 crate::mm::memory::MemoryObject;
6use crate::task::ThreadLockupDetector;
7use crate::task::dynamic_thread_spawner::{DynamicThreadSpawner, SpawnRequestBuilder};
8use starnix_logging::log_debug;
9use starnix_sync::{CondVar, LockDepMutex, PrefetchQueueStateLock, PrefetchWorkerReceiversLock};
10use static_assertions::const_assert;
11use std::collections::VecDeque;
12use std::sync::Arc;
13use std::sync::mpsc::{Receiver, sync_channel};
14
15/// Default maximum number of bytes of requests waiting in the queue at a time (100 MiB).
16const DEFAULT_PREFETCH_QUEUE_MAX_BYTES: u64 = 100 * 1024 * 1024;
17const WORKER_THREAD_COUNT: usize = 2;
18
19#[derive(Debug)]
20pub struct PrefetchMessage {
21    pub memory: Arc<MemoryObject>,
22    pub offset: u64,
23    pub length: u64,
24}
25
26struct QueueState {
27    queue: VecDeque<PrefetchMessage>,
28    queued_bytes: u64,
29    shutdown: bool,
30}
31
32struct PrefetchQueue {
33    state: LockDepMutex<QueueState, PrefetchQueueStateLock>,
34    condvar: CondVar,
35    max_bytes: u64,
36}
37
38impl PrefetchQueue {
39    fn new(max_bytes: u64) -> Self {
40        Self {
41            state: LockDepMutex::new(QueueState {
42                queue: VecDeque::new(),
43                queued_bytes: 0,
44                shutdown: false,
45            }),
46            condvar: CondVar::new(),
47            max_bytes,
48        }
49    }
50
51    fn enqueue(&self, msg: PrefetchMessage) -> bool {
52        if msg.length > self.max_bytes {
53            return false;
54        }
55        {
56            let mut state = self.state.lock();
57            if state.shutdown {
58                return false;
59            }
60            // Evict the oldest requests until there is enough quota for the new request. Don't
61            // worry about overflow here as the max message length above should always be small
62            // enough to avoid it. See const_assert in `PrefetchThreadPool::new()`.
63            while state.queued_bytes + msg.length > self.max_bytes {
64                if let Some(oldest) = state.queue.pop_front() {
65                    state.queued_bytes -= oldest.length;
66                } else {
67                    break;
68                }
69            }
70            state.queued_bytes += msg.length;
71            state.queue.push_back(msg);
72        }
73        self.condvar.notify_one();
74        true
75    }
76}
77
78fn worker_loop(queue: &Arc<PrefetchQueue>) {
79    loop {
80        let msg = {
81            let mut state = queue.state.lock();
82            while state.queue.is_empty() && !state.shutdown {
83                let _waiting_guard = ThreadLockupDetector::pause_tracking();
84                queue.condvar.wait(&mut state);
85            }
86            if state.shutdown {
87                break;
88            }
89            let msg = state.queue.pop_front().unwrap();
90            state.queued_bytes -= msg.length;
91            msg
92        };
93
94        if let Err(status) = msg.memory.op_range(zx::VmoOp::PREFETCH, msg.offset, msg.length) {
95            log_debug!(status:%; "prefetch in background thread failed");
96        }
97    }
98}
99
100/// A thread pool for handling prefetch requests in the background. Has a maximum queue size where
101/// messages that are too large are just dropped and when the queue has too many pending bytes to
102/// service it drops the oldest request first since this is meant to support madvise where the
103/// requests are optional, and often time sensitive.
104pub struct PrefetchThreadPool {
105    queue: Arc<PrefetchQueue>,
106    worker_receivers: LockDepMutex<Vec<Receiver<()>>, PrefetchWorkerReceiversLock>,
107}
108
109impl PrefetchThreadPool {
110    pub fn new(spawner: &DynamicThreadSpawner) -> Self {
111        // Makes it impossible to overflow when checking for space so we can skip all the
112        // verifications.
113        const_assert!(DEFAULT_PREFETCH_QUEUE_MAX_BYTES < u64::MAX / 2);
114        let queue = Arc::new(PrefetchQueue::new(DEFAULT_PREFETCH_QUEUE_MAX_BYTES));
115        let mut worker_receivers = Vec::with_capacity(WORKER_THREAD_COUNT);
116        for _ in 0..WORKER_THREAD_COUNT {
117            let queue_clone = queue.clone();
118            let (sender, receiver) = sync_channel(1);
119            let req = SpawnRequestBuilder::new()
120                .with_debug_name("kthread-prefetch")
121                .with_sync_closure(move |_current_task| {
122                    worker_loop(&queue_clone);
123                    let _ = sender.send(());
124                })
125                .build();
126            spawner.spawn_from_request(req);
127            worker_receivers.push(receiver);
128        }
129        Self { queue, worker_receivers: LockDepMutex::new(worker_receivers) }
130    }
131
132    pub fn enqueue(&self, msg: PrefetchMessage) -> bool {
133        self.queue.enqueue(msg)
134    }
135}
136
137impl Drop for PrefetchThreadPool {
138    fn drop(&mut self) {
139        {
140            let mut state = self.queue.state.lock();
141            state.shutdown = true;
142            state.queued_bytes = 0;
143            state.queue.clear();
144        }
145        self.queue.condvar.notify_all();
146        let receivers = std::mem::take(&mut *self.worker_receivers.lock());
147        for receiver in receivers {
148            let _ = receiver.recv();
149        }
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156    use crate::testing::spawn_kernel_and_run;
157    use assert_matches::assert_matches;
158
159    #[::fuchsia::test]
160    async fn test_prefetch_thread_pool_execution() {
161        spawn_kernel_and_run(async |current_task| {
162            let pool = PrefetchThreadPool::new(current_task.kernel().kthreads.spawner());
163            let page_size = zx::system_get_page_size() as u64;
164
165            let pager = zx::Pager::create(zx::PagerOptions::empty()).unwrap();
166            let port = zx::Port::create();
167            let vmo = pager.create_vmo(zx::VmoOptions::empty(), &port, 1, page_size).unwrap();
168            let child_vmo = vmo.create_child(zx::VmoChildOptions::REFERENCE, 0, 0).unwrap();
169            let memory = Arc::new(MemoryObject::from(child_vmo));
170
171            let enqueued = pool.enqueue(PrefetchMessage { memory, offset: 0, length: page_size });
172            assert!(enqueued);
173
174            // Wait for pager request triggered by the worker's PREFETCH.
175            let packet = port.wait(zx::MonotonicInstant::INFINITE).unwrap();
176            assert_eq!(packet.key(), 1);
177
178            let aux_vmo = zx::Vmo::create(page_size).unwrap();
179            pager.supply_pages(&vmo, 0..page_size, &aux_vmo, 0).unwrap();
180        })
181        .await;
182    }
183
184    #[::fuchsia::test]
185    fn test_prefetch_queue_eviction() {
186        let page_size = zx::system_get_page_size() as u64;
187        let vmo = zx::Vmo::create(page_size).unwrap();
188        let memory = Arc::new(MemoryObject::from(vmo));
189        let queue = PrefetchQueue::new(100);
190
191        // A single request exceeding the quota must be dropped.
192        assert!(!queue.enqueue(PrefetchMessage { memory: memory.clone(), offset: 0, length: 101 }));
193        {
194            let state = queue.state.lock();
195            assert_eq!(state.queued_bytes, 0);
196            assert!(state.queue.is_empty());
197        }
198
199        // Enqueueing requests within quota succeeds without eviction.
200        assert!(queue.enqueue(PrefetchMessage { memory: memory.clone(), offset: 10, length: 40 }));
201        assert!(queue.enqueue(PrefetchMessage { memory: memory.clone(), offset: 20, length: 40 }));
202        {
203            let state = queue.state.lock();
204            assert_eq!(state.queued_bytes, 80);
205            assert_eq!(state.queue.len(), 2);
206            assert_matches!(&state.queue[0], PrefetchMessage { offset: 10, .. });
207            assert_matches!(&state.queue[1], PrefetchMessage { offset: 20, .. });
208        }
209
210        // Adding a 40-byte request exceeds 100 bytes, evicting the oldest (offset 10).
211        assert!(queue.enqueue(PrefetchMessage { memory: memory.clone(), offset: 30, length: 40 }));
212        {
213            let state = queue.state.lock();
214            assert_eq!(state.queued_bytes, 80);
215            assert_eq!(state.queue.len(), 2);
216            assert_matches!(&state.queue[0], PrefetchMessage { offset: 20, .. });
217            assert_matches!(&state.queue[1], PrefetchMessage { offset: 30, .. });
218        }
219
220        // Adding a 90-byte request evicts multiple older requests until it fits.
221        assert!(queue.enqueue(PrefetchMessage { memory: memory.clone(), offset: 50, length: 90 }));
222        {
223            let state = queue.state.lock();
224            assert_eq!(state.queued_bytes, 90);
225            assert_eq!(state.queue.len(), 1);
226            assert_matches!(&state.queue[0], PrefetchMessage { offset: 50, .. });
227        }
228    }
229
230    #[::fuchsia::test]
231    async fn test_drop_pool_with_outstanding_messages() {
232        spawn_kernel_and_run(async |current_task| {
233            let pool = PrefetchThreadPool::new(current_task.kernel().kthreads.spawner());
234            let page_size = zx::system_get_page_size() as u64;
235            let vmo_size = page_size * (WORKER_THREAD_COUNT as u64 + 1);
236
237            // Create a pager-backed vmo, so that we can control the page-in threads.
238            let pager = zx::Pager::create(zx::PagerOptions::empty()).unwrap();
239            let port = zx::Port::create();
240            let vmo = pager.create_vmo(zx::VmoOptions::empty(), &port, 1, vmo_size).unwrap();
241            let child_vmo = vmo.create_child(zx::VmoChildOptions::REFERENCE, 0, 0).unwrap();
242            let memory = Arc::new(MemoryObject::from(child_vmo));
243
244            // Queue one message for each worker, and one that will never be serviced.
245            for offset in
246                (0..(page_size * (WORKER_THREAD_COUNT as u64 + 1))).step_by(page_size as usize)
247            {
248                pool.enqueue(PrefetchMessage { memory: memory.clone(), offset, length: page_size });
249            }
250
251            // Wait for the requests to be pulled off the queue. Now all workers are blocked.
252            for _ in 0..WORKER_THREAD_COUNT {
253                port.wait(zx::MonotonicInstant::INFINITE).unwrap();
254            }
255
256            // Dropping the pool in the background will prevent any more messages from being
257            // handled by the workers as shutdown begins.
258            let queue = pool.queue.clone();
259            let join_handle = std::thread::spawn(move || {
260                drop(pool);
261            });
262
263            // Wait for the other thread to actually commence shutdown.
264            {
265                let mut state = queue.state.lock();
266                while !state.shutdown {
267                    queue.condvar.wait(&mut state);
268                }
269                assert!(state.queue.is_empty());
270                assert_eq!(state.queued_bytes, 0);
271            }
272
273            // Respond to the two outstanding requests. This allows the drop the complete.
274            let aux_vmo = zx::Vmo::create(page_size).unwrap();
275            for offset in (0..(page_size * WORKER_THREAD_COUNT as u64)).step_by(page_size as usize)
276            {
277                pager.supply_pages(&vmo, offset..(offset + page_size), &aux_vmo, 0).unwrap();
278            }
279
280            // Now the drop will complete.
281            join_handle.join().unwrap();
282
283            // The last request was never received.
284            assert!(port.wait(zx::MonotonicInstant::INFINITE_PAST).is_err());
285        })
286        .await;
287    }
288}
starnix_core/mm/prefetch_pool.rs

starnix_core/mm/
prefetch_pool.rs
