Skip to main content

asc_engine/
prebuild.rs

1//! Building the locator tables of an opened APK ahead of its first query
2//! ([`Engine::prebuild`]), on background threads that give way to queries.
3//!
4//! * Only entries already held by the caches are prepared (the memory cache
5//!   after a warm open, L1, or stored entries, which only need mapping): a
6//!   prebuild never inflates, so on a budget too small for the APK it does
7//!   nothing rather than thrash the LRU.
8//! * Tables persisted in L2 are mapped instead of built; tables built here
9//!   are queued for L2 like a query's.
10//! * Queries come first: the workers run at the lowest CPU priority and
11//!   stop between tables while a search of this engine is running. A query
12//!   needing a table a worker is half way through waits for that build
13//!   (`SearchState` builds each table once) instead of starting over.
14//! * [`Engine::cancel_prebuild`] stops it between tables and waits for the
15//!   workers; it is permanent for the engine (an APK closed for good).
16
17use crate::{Engine, Job, Loaded, Result, WriteJob, lookup_cached, recharge, tiers};
18use asc_cache::Tier;
19use asc_dex::Dex;
20use asc_findrefs::FindRefs;
21use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
22use std::sync::{Arc, Condvar, Mutex};
23use std::time::{Duration, Instant};
24
25/// How a paused worker polls for the searches it gave way to.
26const YIELD_POLL: Duration = Duration::from_millis(1);
27/// L2 writes a prebuild lets queue up before it waits for the writer.
28const MAX_PENDING_WRITES: usize = 2;
29
30/// The tables a prebuild prepares, in order (every search needs the
31/// instruction table; strings and types back most member queries too). The
32/// call graph is left to the first walk: it is large and rarely needed.
33const TABLES: [fn(&FindRefs<'_, '_, '_>) -> asc_findrefs::Result<()>; 5] = [
34    |f| f.insn_locator().map(drop),
35    |f| f.string_locator().map(drop),
36    |f| f.type_locator().map(drop),
37    |f| f.method_locator().map(drop),
38    |f| f.field_locator().map(drop),
39];
40
41/// Background build state of one engine.
42#[derive(Default)]
43pub(crate) struct Prebuild {
44    cancel: AtomicBool,
45    running: Mutex<bool>,
46    done: Condvar,
47    /// Searches of this engine in progress: workers wait while any runs.
48    queries: AtomicUsize,
49}
50
51/// Held by a search for its duration ([`Prebuild::query`]).
52pub(crate) struct QueryGuard<'p>(&'p Prebuild);
53
54impl Drop for QueryGuard<'_> {
55    fn drop(&mut self) {
56        self.0.queries.fetch_sub(1, Ordering::Relaxed);
57    }
58}
59
60impl Prebuild {
61    /// Mark a search as running until the guard drops.
62    pub(crate) fn query(&self) -> QueryGuard<'_> {
63        self.queries.fetch_add(1, Ordering::Relaxed);
64        QueryGuard(self)
65    }
66
67    pub(crate) fn cancel(&self) {
68        self.cancel.store(true, Ordering::Relaxed);
69    }
70
71    fn cancelled(&self) -> bool {
72        self.cancel.load(Ordering::Relaxed)
73    }
74
75    /// Wait while a search runs; `false` once cancelled.
76    fn give_way(&self) -> bool {
77        while self.queries.load(Ordering::Relaxed) > 0 && !self.cancelled() {
78            std::thread::sleep(YIELD_POLL);
79        }
80        !self.cancelled()
81    }
82}
83
84/// What [`Engine::prebuild`] did.
85#[derive(Debug, Clone, Default, PartialEq, Eq)]
86pub struct PrebuildReport {
87    /// Logical DEX files prepared.
88    pub dex: usize,
89    /// Of those, the ones whose tables came from L2.
90    pub from_disk: usize,
91    /// Root entries skipped because no cache held them.
92    pub skipped: usize,
93    /// Stopped by [`Engine::cancel_prebuild`] (or already running).
94    pub cancelled: bool,
95    pub elapsed_us: u128,
96}
97
98impl Engine {
99    /// Build (or map from L2) the locator tables of every DEX entry the
100    /// caches hold, so the first query of each kind skips its table builds.
101    /// Blocks until done; meant for a background thread after a warm open.
102    /// Does nothing without a memory cache (the tables would be dropped).
103    pub fn prebuild(&self) -> Result<PrebuildReport> {
104        let t0 = Instant::now();
105        let mut report = PrebuildReport::default();
106        let Some(cache) = self.cache.as_deref().filter(|c| c.memory.is_some()) else { return Ok(report) };
107        {
108            let mut running = self.prebuild.running.lock().unwrap();
109            if *running || self.prebuild.cancelled() {
110                report.cancelled = true;
111                return Ok(report);
112            }
113            *running = true;
114        }
115        let _done = Finished(&self.prebuild);
116        let apk = self.open()?;
117        let mut entries = apk.dex_entries()?;
118        // Largest first, as a search runs them.
119        entries.reverse();
120        let key = self.apk_key(&apk);
121        let l2 = cache.disk.as_deref().filter(|d| d.uses(Tier::Derived)).zip(key.as_ref());
122        let next = AtomicUsize::new(0);
123        let tally = Mutex::new(&mut report);
124        std::thread::scope(|s| {
125            for _ in 0..self.threads().min(entries.len()) {
126                s.spawn(|| {
127                    lower_priority();
128                    while self.prebuild.give_way() {
129                        let Some(entry) = entries.get(next.fetch_add(1, Ordering::Relaxed)) else { break };
130                        let cached = match lookup_cached(cache, key.as_ref(), entry, &self.stats) {
131                            Some((e, _)) => Some(e),
132                            // Stored entries are only mapped: as cheap as a lookup.
133                            None if entry.method == asc_apk::METHOD_STORED => {
134                                match crate::load_entry(&apk, key.as_ref(), entry, Some(cache), &self.stats, Some(&self.prebuild.cancel)) {
135                                    Ok((Loaded::Cached(e), _)) => Some(e),
136                                    _ => None,
137                                }
138                            }
139                            None => None,
140                        };
141                        let Some(e) = cached else {
142                            tally.lock().unwrap().skipped += 1;
143                            continue;
144                        };
145                        for (i, l) in e.logical().iter().enumerate() {
146                            let Ok(dex) = Dex::with_name(&l.data, &l.name) else { continue };
147                            let mapped = l2.is_some_and(|(disk, k)| tiers::load_derived(disk, k, l, &dex));
148                            let (mask, t) = (l.state.table_mask(), Instant::now());
149                            let fr = FindRefs::with_state(&dex, &l.state);
150                            for build in TABLES {
151                                // A damaged DEX fails its query later with the same error.
152                                if !self.prebuild.give_way() || build(&fr).is_err() {
153                                    break;
154                                }
155                            }
156                            if let Some((disk, k)) = l2
157                                && let Some(job) = tiers::derived_job(disk, &e, i, mask, t.elapsed())
158                            {
159                                cache.put_async(WriteJob { key: k.clone(), apk: Arc::from(apk.path()), kind: Job::Derived(job) });
160                                // Built tables stay on the heap until the writer maps them:
161                                // building far ahead of it only adds up heap.
162                                while cache.derived_pending.load(Ordering::Relaxed) > MAX_PENDING_WRITES && !self.prebuild.cancelled() {
163                                    std::thread::sleep(YIELD_POLL);
164                                }
165                            }
166                            let mut tally = tally.lock().unwrap();
167                            tally.dex += 1;
168                            tally.from_disk += mapped as usize;
169                        }
170                        recharge(Some(cache), key.as_ref(), &e);
171                    }
172                });
173            }
174        });
175        report.cancelled = self.prebuild.cancelled();
176        report.elapsed_us = t0.elapsed().as_micros();
177        self.log(format!(
178            "[APK] prebuild dex={} from_disk={} skipped={} cancelled={} total={} us",
179            report.dex, report.from_disk, report.skipped, report.cancelled, report.elapsed_us
180        ));
181        Ok(report)
182    }
183
184    /// Stop a running [`Engine::prebuild`] (and any later one) and wait for
185    /// its workers to finish the table they are on.
186    pub fn cancel_prebuild(&self) {
187        self.prebuild.cancel();
188        let mut running = self.prebuild.running.lock().unwrap();
189        while *running {
190            running = self.prebuild.done.wait(running).unwrap();
191        }
192    }
193}
194
195/// Clears `running` when a prebuild ends, however it ends.
196struct Finished<'p>(&'p Prebuild);
197
198impl Drop for Finished<'_> {
199    fn drop(&mut self) {
200        *self.0.running.lock().unwrap_or_else(std::sync::PoisonError::into_inner) = false;
201        self.0.done.notify_all();
202    }
203}
204
205/// Run the calling thread at the lowest CPU priority (nice 19 on Linux, the
206/// utility QoS class on Apple platforms), so background work (prebuilds,
207/// cache writers) only uses cycles queries leave idle. Utility rather than
208/// background QoS: the latter also throttles the thread's disk I/O.
209pub(crate) fn lower_priority() {
210    #[cfg(target_vendor = "apple")]
211    {
212        const QOS_CLASS_UTILITY: std::os::raw::c_uint = 0x11;
213        unsafe extern "C" {
214            fn pthread_set_qos_class_self_np(
215                qos_class: std::os::raw::c_uint,
216                relative_priority: std::os::raw::c_int,
217            ) -> std::os::raw::c_int;
218        }
219        // SAFETY: plain call on the calling thread; it only changes that
220        // thread's scheduling class.
221        unsafe {
222            pthread_set_qos_class_self_np(QOS_CLASS_UTILITY, 0);
223        }
224    }
225    #[cfg(all(target_os = "linux", target_env = "gnu"))]
226    {
227        unsafe extern "C" {
228            fn gettid() -> std::os::raw::c_int;
229            fn setpriority(which: std::os::raw::c_int, who: std::os::raw::c_uint, prio: std::os::raw::c_int) -> std::os::raw::c_int;
230        }
231        const PRIO_PROCESS: std::os::raw::c_int = 0;
232        // SAFETY: plain syscalls on the calling thread's own id; Linux applies
233        // the nice value to that thread only.
234        unsafe {
235            setpriority(PRIO_PROCESS, gettid() as std::os::raw::c_uint, 19);
236        }
237    }
238}