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::{LockExt, 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.
33pub(crate) const 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)]
86#[non_exhaustive]
87pub struct PrebuildReport {
88    /// Logical DEX files prepared.
89    pub dex: usize,
90    /// Of those, the ones whose tables came from L2.
91    pub from_disk: usize,
92    /// Root entries skipped because no cache held them.
93    pub skipped: usize,
94    /// Stopped by [`Engine::cancel_prebuild`] (or already running).
95    pub cancelled: bool,
96    /// Wall-clock microseconds of the run.
97    pub elapsed_us: u128,
98}
99
100impl Engine {
101    /// Build (or map from L2) the locator tables of every DEX entry the
102    /// caches hold, so the first query of each kind skips its table builds.
103    /// Blocks until done; meant for a background thread after a warm open.
104    /// Does nothing without a memory cache (the tables would be dropped).
105    ///
106    /// # Errors
107    ///
108    /// [`EngineError::Apk`](crate::EngineError::Apk) when the APK cannot be
109    /// opened. A DEX that fails to parse or build tables is skipped (its
110    /// queries report the error later).
111    ///
112    /// # Panics
113    ///
114    /// When a worker thread panics (a bug); the panic is propagated.
115    pub fn prebuild(&self) -> Result<PrebuildReport> {
116        let t0 = Instant::now();
117        let mut report = PrebuildReport::default();
118        let Some(cache) = self.cache.as_deref().filter(|c| c.memory.is_some()) else { return Ok(report) };
119        {
120            let mut running = self.prebuild.running.lock_or_default();
121            if *running || self.prebuild.cancelled() {
122                report.cancelled = true;
123                return Ok(report);
124            }
125            *running = true;
126        }
127        let _done = Finished(&self.prebuild);
128        let apk = self.open()?;
129        let mut entries = apk.dex_entries()?;
130        // Largest first, as a search runs them.
131        entries.reverse();
132        let key = self.apk_key(&apk);
133        let l2 = cache.disk.as_deref().filter(|d| d.uses(Tier::Derived)).zip(key.as_ref());
134        let next = AtomicUsize::new(0);
135        let tally = Mutex::new(&mut report);
136        std::thread::scope(|s| {
137            for _ in 0..self.threads().min(entries.len()) {
138                s.spawn(|| {
139                    lower_priority();
140                    while self.prebuild.give_way() {
141                        let Some(entry) = entries.get(next.fetch_add(1, Ordering::Relaxed)) else { break };
142                        let cached = match lookup_cached(cache, key.as_ref(), entry, &self.stats) {
143                            Some((e, _)) => Some(e),
144                            // Stored entries are only mapped: as cheap as a lookup.
145                            None if entry.method == asc_apk::METHOD_STORED => {
146                                match crate::load_entry(&apk, key.as_ref(), entry, Some(cache), &self.stats, Some(&self.prebuild.cancel)) {
147                                    Ok((Loaded::Cached(e), _)) => Some(e),
148                                    _ => None,
149                                }
150                            }
151                            None => None,
152                        };
153                        let Some(e) = cached else {
154                            tally.lock().expect("prebuild tally mutex poisoned").skipped += 1;
155                            continue;
156                        };
157                        for (i, l) in e.logical().iter().enumerate() {
158                            let Ok(dex) = Dex::with_name(&l.data, &l.name) else { continue };
159                            let mapped = l2.is_some_and(|(disk, k)| tiers::load_derived(disk, k, l, &dex));
160                            let (mask, t) = (l.state.table_mask(), Instant::now());
161                            let fr = FindRefs::with_state(&dex, &l.state);
162                            for build in TABLES {
163                                // A damaged DEX fails its query later with the same error.
164                                if !self.prebuild.give_way() || build(&fr).is_err() {
165                                    break;
166                                }
167                            }
168                            if let Some((disk, k)) = l2
169                                && let Some(job) = tiers::derived_job(disk, &e, i, mask, t.elapsed())
170                            {
171                                cache.put_async(WriteJob { key: k.clone(), apk: Arc::from(apk.path()), kind: Job::Derived(job) });
172                                // Built tables stay on the heap until the writer maps them:
173                                // building far ahead of it only adds up heap.
174                                while cache.derived_pending.load(Ordering::Relaxed) > MAX_PENDING_WRITES && !self.prebuild.cancelled() {
175                                    std::thread::sleep(YIELD_POLL);
176                                }
177                            }
178                            let mut tally = tally.lock().expect("prebuild tally mutex poisoned");
179                            tally.dex += 1;
180                            tally.from_disk += mapped as usize;
181                        }
182                        recharge(Some(cache), key.as_ref(), &e);
183                    }
184                });
185            }
186        });
187        report.cancelled = self.prebuild.cancelled();
188        report.elapsed_us = t0.elapsed().as_micros();
189        self.log(format!(
190            "[APK] prebuild dex={} from_disk={} skipped={} cancelled={} total={} us",
191            report.dex, report.from_disk, report.skipped, report.cancelled, report.elapsed_us
192        ));
193        Ok(report)
194    }
195
196    /// Stop a running [`Engine::prebuild`] (and any later one) and wait for
197    /// its workers to finish the table they are on.
198    ///
199    /// # Panics
200    ///
201    /// When the prebuild's state mutex is poisoned (a prebuild thread
202    /// panicked while holding it).
203    pub fn cancel_prebuild(&self) {
204        self.prebuild.cancel();
205        let mut running = self.prebuild.running.lock_or_default();
206        while *running {
207            running = self.prebuild.done.wait(running).expect("prebuild flag mutex poisoned");
208        }
209    }
210}
211
212/// Clears `running` when a prebuild ends, however it ends.
213struct Finished<'p>(&'p Prebuild);
214
215impl Drop for Finished<'_> {
216    fn drop(&mut self) {
217        *self.0.running.lock_or_default() = false;
218        self.0.done.notify_all();
219    }
220}
221
222/// Run the calling thread at the lowest CPU priority (nice 19 on Linux, the
223/// utility QoS class on Apple platforms), so background work (prebuilds,
224/// cache writers) only uses cycles queries leave idle. Utility rather than
225/// background QoS: the latter also throttles the thread's disk I/O.
226pub(crate) fn lower_priority() {
227    #[cfg(target_vendor = "apple")]
228    {
229        const QOS_CLASS_UTILITY: std::os::raw::c_uint = 0x11;
230        unsafe extern "C" {
231            fn pthread_set_qos_class_self_np(
232                qos_class: std::os::raw::c_uint,
233                relative_priority: std::os::raw::c_int,
234            ) -> std::os::raw::c_int;
235        }
236        // SAFETY: plain call on the calling thread; it only changes that
237        // thread's scheduling class.
238        unsafe {
239            pthread_set_qos_class_self_np(QOS_CLASS_UTILITY, 0);
240        }
241    }
242    #[cfg(all(target_os = "linux", target_env = "gnu"))]
243    {
244        unsafe extern "C" {
245            fn gettid() -> std::os::raw::c_int;
246            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;
247        }
248        const PRIO_PROCESS: std::os::raw::c_int = 0;
249        // SAFETY: plain syscalls on the calling thread's own id; Linux applies
250        // the nice value to that thread only.
251        unsafe {
252            setpriority(PRIO_PROCESS, gettid() as std::os::raw::c_uint, 19);
253        }
254    }
255}