asc_engine/
tiers.rs

1//! The engine's use of the disk cache tiers beyond inflated DEX entries, and
2//! the memory budget shared by the resident hosts of one user.
3//!
4//! * L2 (`asc_cache::Tier::Derived`): the locator tables of each logical DEX
5//!   ([`LOCATORS_KIND`]). A search first adopts the tables a previous
6//!   process persisted (mapped, so every process shares one copy of their
7//!   pages), and queues the tables it had to build for writing once they took
8//!   long enough to build (`DiskCache::admits_derived`). Once written, the copy in the memory cache is
9//!   swapped for a view of the file, like inflated bytes are remapped to L1.
10//! * L3 (`asc_cache::Tier::Rendered`): the result rows of a `find_refs`
11//!   query, replayed by a later run of the same query on the same APK with
12//!   the same build of the tool.
13//!
14//! * L2 also keeps the class graph (`hierarchy::GRAPH_KIND`, one per APK),
15//!   used in place where it is mapped.
16//!
17//! Other derived indexes (call graph edges, resource tables, manifest
18//! summaries) go to L2 the same way: serialise them into bytes that can be
19//! viewed in place, store them with `DiskCache::put_derived(key, name, kind,
20//! version, build_cost, bytes)` from the cache writer, and map them back with
21//! `DiskCache::get_derived(key, name, kind, version)`. Pick a `kind` of their
22//! own and bump `version` whenever the layout changes; per-APK indexes use a
23//! fixed `name` such as `apk`.
24
25use crate::{
26    CachedEntry, DexBytes, FindOptions, LogicalDex, Matched, MemoryPolicy, RefBatch, RefResult, RefSite, Shape, SharedMemory, TextRows,
27    available_memory,
28};
29use asc_cache::{ApkKey, DiskCache, ProcessRegistry};
30use asc_dex::Dex;
31use asc_findrefs::{ClassRelation, Query, SearchState};
32use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
33use std::sync::{Arc, Mutex, OnceLock};
34use std::time::{Duration, Instant};
35
36/// L2 kind of the persisted locator tables.
37pub(crate) const LOCATORS_KIND: &str = "locators";
38/// `find_refs` results larger than this (encoded) are not kept in L3.
39pub(crate) const L3_MAX_RESULT: usize = 4 << 20;
40/// Text rows ([`Shape::Text`]) larger than this are not kept in L3, nor
41/// beyond an eighth of the disk cache: a few broad queries must not evict the
42/// inflated DEX files every query needs.
43const L3_MAX_TEXT: usize = 32 << 20;
44/// Layout version of the encoded `find_refs` results.
45const RESULT_FORMAT: u32 = 6;
46/// Ends the encoded batches (where an entry name's length would be), so a
47/// truncated file is told from a complete one; the run's notes follow it.
48const RESULT_END: u32 = u32::MAX;
49
50/// L2 bookkeeping of one logical DEX, shared by the entries that replace
51/// each other in the memory cache.
52#[derive(Default)]
53pub(crate) struct Persisted {
54    /// L2 has been consulted for this DEX.
55    tried: AtomicBool,
56    /// Tables on disk or queued for writing (`SearchState::table_mask` bits).
57    queued: AtomicU32,
58    /// Time spent building the tables so far.
59    cost_us: AtomicU64,
60}
61
62/// Adopt the locator tables persisted for `l` (once per logical DEX).
63pub(crate) fn load_derived(disk: &DiskCache, key: &ApkKey, l: &LogicalDex, dex: &Dex<'_>) -> bool {
64    if l.persisted.tried.swap(true, Ordering::Relaxed) {
65        return false;
66    }
67    let Some(d) = disk.get_derived(key, &l.name, LOCATORS_KIND, asc_findrefs::PERSIST_FORMAT_VERSION) else { return false };
68    let cost = d.build_cost.as_micros() as u64;
69    let Some(state) = SearchState::from_persisted(dex, Arc::new(d)) else { return false };
70    l.persisted.queued.fetch_or(state.table_mask(), Ordering::Relaxed);
71    l.persisted.cost_us.fetch_add(cost, Ordering::Relaxed);
72    l.state.fill_from(state);
73    true
74}
75
76/// Tables of a logical DEX to persist.
77pub(crate) struct DerivedJob {
78    /// Memory cache entry and logical index the tables belong to.
79    entry: String,
80    logical: usize,
81    name: String,
82    data: DexBytes,
83    state: Arc<SearchState>,
84    persisted: Arc<Persisted>,
85    /// The `queued` bits this job claimed.
86    claimed: u32,
87}
88
89impl DerivedJob {
90    /// Heap the job keeps alive until written: the tables it persists.
91    pub(crate) fn held_bytes(&self) -> usize {
92        self.state.heap_bytes()
93    }
94
95    /// Not written after all: a later search may queue these tables again.
96    pub(crate) fn abandon(self) {
97        self.persisted.queued.fetch_and(!self.claimed, Ordering::Relaxed);
98    }
99}
100
101/// A write of the tables of `entry.logical()[idx]`, when a search that
102/// started with `mask_before` built tables worth persisting (`spent` is
103/// what it took).
104pub(crate) fn derived_job(disk: &DiskCache, entry: &CachedEntry, idx: usize, mask_before: u32, spent: Duration) -> Option<DerivedJob> {
105    let l = entry.logical().get(idx)?;
106    let mask = l.state.table_mask();
107    let p = &l.persisted;
108    if mask == mask_before && mask & !p.queued.load(Ordering::Relaxed) == 0 {
109        return None;
110    }
111    let cost = if mask != mask_before {
112        p.cost_us.fetch_add(spent.as_micros() as u64, Ordering::Relaxed) + spent.as_micros() as u64
113    } else {
114        p.cost_us.load(Ordering::Relaxed)
115    };
116    if !disk.admits_derived(Duration::from_micros(cost)) {
117        return None;
118    }
119    let claimed = mask & !p.queued.fetch_or(mask, Ordering::Relaxed);
120    if claimed == 0 {
121        return None;
122    }
123    Some(DerivedJob {
124        entry: entry.name.clone(),
125        logical: idx,
126        name: l.name.clone(),
127        data: Arc::clone(&l.data),
128        state: Arc::clone(&l.state),
129        persisted: Arc::clone(&l.persisted),
130        claimed,
131    })
132}
133
134/// Write a [`DerivedJob`] (on the cache writer thread), then swap the heap
135/// tables of the memory cache's entry for views of the file.
136pub(crate) fn write_derived(disk: &DiskCache, memory: Option<&SharedMemory>, key: &ApkKey, job: DerivedJob) {
137    let Ok(dex) = Dex::with_name(&job.data, &job.name) else { return };
138    let mask = job.state.table_mask();
139    let bytes = job.state.persist(&dex);
140    let cost = Duration::from_micros(job.persisted.cost_us.load(Ordering::Relaxed));
141    let version = asc_findrefs::PERSIST_FORMAT_VERSION;
142    if disk.put_derived(key, &job.name, LOCATORS_KIND, version, cost, &bytes).is_err() {
143        return;
144    }
145    drop(bytes);
146    let Some(memory) = memory else { return };
147    let Some(mapped) =
148        disk.get_derived(key, &job.name, LOCATORS_KIND, version).and_then(|d| SearchState::from_persisted(&dex, Arc::new(d)))
149    else {
150        return;
151    };
152    // Tables built after the snapshot would be lost by the swap: keep the
153    // heap state then (the next write maps them all).
154    if mapped.table_mask() != mask || job.state.table_mask() != mask {
155        return;
156    }
157    let mut mem = memory.lock().unwrap();
158    let Some(current) = mem.peek(key, &job.entry) else { return };
159    let still_ours = current.logical.get().and_then(|l| l.get(job.logical)).is_some_and(|l| Arc::ptr_eq(&l.state, &job.state));
160    if still_ours {
161        let new = Arc::new(current.with_state(job.logical, Arc::new(mapped)));
162        let size = new.charged_bytes();
163        if mem.replace(key, &job.entry, &current, Arc::clone(&new)) {
164            mem.resize(key, &job.entry, &new, size);
165        }
166    }
167}
168
169/// Fingerprint of this build of the tool: the crate version plus the size
170/// and modification time of the running executable, so a rebuilt binary
171/// never replays results of the old one.
172fn tool_version() -> &'static str {
173    static VERSION: OnceLock<String> = OnceLock::new();
174    VERSION.get_or_init(|| {
175        let exe = std::env::current_exe().ok().and_then(|p| std::fs::metadata(p).ok());
176        let stamp = exe.map(|m| {
177            let mtime = m.modified().ok().and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()).map_or(0, |d| d.as_nanos());
178            format!("{}.{mtime}", m.len())
179        });
180        format!("{}+{}", env!("CARGO_PKG_VERSION"), stamp.unwrap_or_default())
181    })
182}
183
184/// The L3 key of a `find_refs` query.
185pub(crate) fn findrefs_key(query: &Query, opts: FindOptions, shape: Shape) -> String {
186    let query = match shape {
187        Shape::Results => format!("{query:?}\0{opts:?}"),
188        Shape::Text => format!("{query:?}\0{opts:?}\0text"),
189    };
190    asc_cache::rendered_key("findrefs", tool_version(), &query)
191}
192
193/// The most bytes a `find_refs` result of `shape` may take in `disk`'s L3.
194pub(crate) fn result_limit(disk: &DiskCache, shape: Shape) -> usize {
195    match shape {
196        Shape::Results => L3_MAX_RESULT,
197        Shape::Text => usize::try_from(disk.capacity_for(0) / 8).unwrap_or(usize::MAX).min(L3_MAX_TEXT),
198    }
199}
200
201/// Encodes the batches of a `find_refs` run for L3 as they stream past.
202/// Runs with an error (an entry that could not be searched may succeed next
203/// time) or results beyond the recorder's limit ([`result_limit`]) are not
204/// kept.
205pub(crate) struct ResultRecorder {
206    out: Vec<u8>,
207    keep: bool,
208    limit: usize,
209}
210
211impl ResultRecorder {
212    pub(crate) fn new(limit: usize) -> ResultRecorder {
213        let mut out = Vec::new();
214        out.extend_from_slice(&RESULT_FORMAT.to_le_bytes());
215        ResultRecorder { out, keep: true, limit }
216    }
217
218    /// Stop recording: the run will not be kept.
219    fn abandon(&mut self) {
220        self.keep = false;
221        self.out = Vec::new();
222    }
223
224    pub(crate) fn record(&mut self, b: &RefBatch) {
225        if !self.keep {
226            return;
227        }
228        // Text rows are the bulk of a broad run: check them before copying.
229        let text_bytes = b.text.as_ref().map_or(0, |t| t.as_str().len() + 4 * t.len());
230        if b.error.is_some() || self.out.len().saturating_add(text_bytes) > self.limit {
231            self.abandon();
232            return;
233        }
234        put_str(&mut self.out, &b.entry_name);
235        put_u32(&mut self.out, b.results.len() as u32);
236        for r in &b.results {
237            put_str(&mut self.out, &r.dex_name);
238            put_u32(&mut self.out, r.method_idx);
239            put_str(&mut self.out, &r.method);
240            put_str(&mut self.out, &r.signature);
241            put_u32(&mut self.out, r.matched.len() as u32);
242            for m in &r.matched {
243                self.out.push(match m {
244                    Matched::String(_) => 0,
245                    Matched::Symbol(_) => 1,
246                    Matched::Field { .. } => 2,
247                });
248                put_str(&mut self.out, m.as_str());
249                if let Matched::Field { access, .. } = m {
250                    self.out.push(match access {
251                        asc_findrefs::FieldAccess::Read => 1,
252                        asc_findrefs::FieldAccess::Write => 2,
253                    });
254                }
255            }
256            self.out.push(match r.relation {
257                None => 0,
258                Some(ClassRelation::Extends) => 1,
259                Some(ClassRelation::Implements) => 2,
260            });
261            put_u32(&mut self.out, r.sites.len() as u32);
262            for site in &r.sites {
263                put_u32(&mut self.out, site.offset);
264                put_u32(&mut self.out, site.matched as u32);
265                self.out.push(match site.field_access {
266                    None => 0,
267                    Some(asc_findrefs::FieldAccess::Read) => 1,
268                    Some(asc_findrefs::FieldAccess::Write) => 2,
269                });
270                put_opt(&mut self.out, site.pc);
271                put_opt(&mut self.out, site.line);
272            }
273            match &r.context {
274                None => self.out.push(0),
275                Some(c) => {
276                    self.out.push(1);
277                    put_str(&mut self.out, c);
278                }
279            }
280        }
281        match &b.text {
282            None => self.out.push(0),
283            Some(t) => {
284                let Ok(ends) = t.ends().iter().map(|&e| u32::try_from(e)).collect::<Result<Vec<_>, _>>() else {
285                    return self.abandon();
286                };
287                self.out.push(1);
288                put_str(&mut self.out, t.as_str());
289                put_u32(&mut self.out, ends.len() as u32);
290                for e in ends {
291                    put_u32(&mut self.out, e);
292                }
293            }
294        }
295        if self.out.len() > self.limit {
296            self.abandon();
297        }
298    }
299
300    /// The encoded run, its `notes` ([`crate::RefSummary::notes`]) last.
301    pub(crate) fn finish(mut self, notes: &[String]) -> Option<Vec<u8>> {
302        put_u32(&mut self.out, RESULT_END);
303        put_u32(&mut self.out, notes.len() as u32);
304        for n in notes {
305            put_str(&mut self.out, n);
306        }
307        self.keep.then_some(self.out)
308    }
309}
310
311fn put_u32(out: &mut Vec<u8>, v: u32) {
312    out.extend_from_slice(&v.to_le_bytes());
313}
314
315fn put_opt(out: &mut Vec<u8>, v: Option<u32>) {
316    match v {
317        None => out.push(0),
318        Some(v) => {
319            out.push(1);
320            put_u32(out, v);
321        }
322    }
323}
324
325fn put_str(out: &mut Vec<u8>, s: &str) {
326    put_u32(out, s.len() as u32);
327    out.extend_from_slice(s.as_bytes());
328}
329
330/// A run [`ResultRecorder`] kept.
331pub(crate) struct Recorded {
332    pub batches: Vec<RefBatch>,
333    pub notes: Vec<String>,
334}
335
336/// The run recorded by [`ResultRecorder`]; `None` for anything malformed.
337/// Text rows stay views of `payload` (typically a mapped L3 file).
338pub(crate) fn decode_batches(payload: impl AsRef<[u8]> + Send + Sync + 'static) -> Option<Recorded> {
339    let payload: Arc<dyn AsRef<[u8]> + Send + Sync> = Arc::new(payload);
340    let bytes = (*payload).as_ref();
341    let mut r = Cursor { bytes, at: 0 };
342    if r.u32()? != RESULT_FORMAT {
343        return None;
344    }
345    let mut out = Vec::new();
346    loop {
347        let len = r.u32()?;
348        if len == RESULT_END {
349            let notes = (0..r.u32()?).map(|_| r.str()).collect::<Option<Vec<_>>>()?;
350            return (r.at == bytes.len()).then_some(Recorded { batches: out, notes });
351        }
352        let entry_name = String::from_utf8(r.take(len as usize)?.to_vec()).ok()?;
353        let n = r.u32()? as usize;
354        let mut results = Vec::with_capacity(n.min(bytes.len()));
355        for _ in 0..n {
356            let dex_name = r.str()?;
357            let method_idx = r.u32()?;
358            let method = r.str()?;
359            let signature = r.str()?;
360            let m = r.u32()? as usize;
361            let matched = (0..m)
362                .map(|_| match r.u8()? {
363                    0 => Some(Matched::String(r.str()?)),
364                    1 => Some(Matched::Symbol(r.str()?)),
365                    2 => {
366                        let symbol = r.str()?;
367                        let access = match r.u8()? {
368                            1 => asc_findrefs::FieldAccess::Read,
369                            2 => asc_findrefs::FieldAccess::Write,
370                            _ => return None,
371                        };
372                        Some(Matched::Field { symbol, access })
373                    }
374                    _ => None,
375                })
376                .collect::<Option<Vec<_>>>()?;
377            let relation = match r.u8()? {
378                0 => None,
379                1 => Some(ClassRelation::Extends),
380                2 => Some(ClassRelation::Implements),
381                _ => return None,
382            };
383            let sites = (0..r.u32()?)
384                .map(|_| {
385                    let offset = r.u32()?;
386                    let matched = r.u32()? as usize;
387                    let field_access = match r.u8()? {
388                        0 => None,
389                        1 => Some(asc_findrefs::FieldAccess::Read),
390                        2 => Some(asc_findrefs::FieldAccess::Write),
391                        _ => return None,
392                    };
393                    Some(RefSite { offset, matched, field_access, pc: r.opt()?, line: r.opt()? })
394                })
395                .collect::<Option<Vec<_>>>()?;
396            let context = match r.u8()? {
397                0 => None,
398                1 => Some(r.str()?),
399                _ => return None,
400            };
401            results.push(RefResult { dex_name, method_idx, method, signature, matched, relation, sites, context });
402        }
403        let text = match r.u8()? {
404            0 => None,
405            1 => {
406                let len = r.u32()? as usize;
407                let start = r.at;
408                r.take(len)?;
409                let ends = (0..r.u32()?).map(|_| r.u32().map(|e| e as usize)).collect::<Option<Vec<_>>>()?;
410                Some(TextRows::shared(Arc::clone(&payload), start..start + len, ends)?)
411            }
412            _ => return None,
413        };
414        out.push(RefBatch { entry_name, results, text, inflate_us: 0, search_us: 0, source: "result", error: None });
415    }
416}
417
418struct Cursor<'b> {
419    bytes: &'b [u8],
420    at: usize,
421}
422
423impl Cursor<'_> {
424    fn take(&mut self, n: usize) -> Option<&[u8]> {
425        let s = self.bytes.get(self.at..self.at.checked_add(n)?)?;
426        self.at += n;
427        Some(s)
428    }
429
430    fn u8(&mut self) -> Option<u8> {
431        Some(self.take(1)?[0])
432    }
433
434    fn u32(&mut self) -> Option<u32> {
435        Some(u32::from_le_bytes(self.take(4)?.try_into().ok()?))
436    }
437
438    /// `Some(None)` / `Some(Some(v))` for a present field, `None` when malformed.
439    fn opt(&mut self) -> Option<Option<u32>> {
440        match self.u8()? {
441            0 => Some(None),
442            1 => Some(Some(self.u32()?)),
443            _ => None,
444        }
445    }
446
447    fn str(&mut self) -> Option<String> {
448        let n = self.u32()? as usize;
449        String::from_utf8(self.take(n)?.to_vec()).ok()
450    }
451}
452
453/// How often a memory cache re-reads the registry and resizes itself.
454const REBALANCE_EVERY: Duration = Duration::from_secs(1);
455
456/// The process's registration: one per process, however many caches it has.
457struct Registration {
458    caches: usize,
459    registry: Option<ProcessRegistry>,
460    /// Bytes held by all of this process's memory caches.
461    held: u64,
462    /// Requests in progress ([`Busy`]).
463    requests: usize,
464}
465
466impl Registration {
467    fn publish(&self) {
468        if let Some(registry) = &self.registry {
469            let _ = registry.update(self.held, self.requests > 0);
470        }
471    }
472}
473
474static REGISTRATION: Mutex<Registration> = Mutex::new(Registration { caches: 0, registry: None, held: 0, requests: 0 });
475
476/// Marks this process busy in the user's process registry while a request
477/// (an MCP tool call, a daemon command) runs, so the others split the CPUs
478/// with it from its first moment ([`auto_threads`]) and take them back as soon
479/// as it ends. Nested and concurrent requests count once.
480#[must_use = "the process is busy only while the guard lives"]
481pub struct Busy(());
482
483impl Busy {
484    pub fn enter() -> Busy {
485        let mut reg = REGISTRATION.lock().unwrap();
486        reg.requests += 1;
487        if reg.requests == 1 {
488            reg.publish();
489        }
490        Busy(())
491    }
492}
493
494impl Drop for Busy {
495    fn drop(&mut self) {
496        let mut reg = REGISTRATION.lock().unwrap();
497        reg.requests -= 1;
498        if reg.requests == 0 {
499            reg.publish();
500        }
501    }
502}
503
504/// The memory budget of one memory cache: published to the user's process
505/// registry and, when automatic, resized to this process's share of one
506/// budget for all registered processes ([`MemoryPolicy::share`]).
507pub(crate) struct MemoryBudget {
508    auto: Option<MemoryPolicy>,
509    last: Mutex<Option<Instant>>,
510    published: AtomicU64,
511}
512
513impl MemoryBudget {
514    pub(crate) fn new(auto: Option<MemoryPolicy>) -> MemoryBudget {
515        let mut reg = REGISTRATION.lock().unwrap();
516        reg.caches += 1;
517        if reg.registry.is_none() {
518            reg.registry = ProcessRegistry::default_dir().and_then(ProcessRegistry::open);
519        }
520        MemoryBudget { auto, last: Mutex::new(None), published: AtomicU64::new(0) }
521    }
522
523    /// Publish what `memory` holds and, when automatic, resize it to this
524    /// process's share. At most once per [`REBALANCE_EVERY`].
525    pub(crate) fn maybe_rebalance(&self, memory: &SharedMemory) {
526        {
527            let mut last = self.last.lock().unwrap();
528            if last.is_some_and(|t| t.elapsed() < REBALANCE_EVERY) {
529                return;
530            }
531            *last = Some(Instant::now());
532        }
533        let used = {
534            let m = memory.lock().unwrap();
535            (m.used_bytes() + m.reserved_bytes()) as u64
536        };
537        let previous = self.published.swap(used, Ordering::Relaxed);
538        let (held, others) = {
539            let mut reg = REGISTRATION.lock().unwrap();
540            reg.held = (reg.held + used).saturating_sub(previous);
541            reg.publish();
542            let Some(registry) = &reg.registry else { return };
543            (reg.held, registry.others())
544        };
545        if let Some(policy) = self.auto {
546            let cap = policy.share(available_memory(), held, &others);
547            memory.lock().unwrap().set_capacity(cap);
548        }
549    }
550}
551
552impl Drop for MemoryBudget {
553    fn drop(&mut self) {
554        let mut reg = REGISTRATION.lock().unwrap();
555        reg.held = reg.held.saturating_sub(self.published.load(Ordering::Relaxed));
556        reg.caches -= 1;
557        if reg.caches == 0 {
558            // Unregisters (removes this process's entry).
559            reg.registry = None;
560        } else {
561            reg.publish();
562        }
563    }
564}
565
566/// Busy for [`auto_threads`], for an entry that does not say ([`Busy`]
567/// came later): one updated this recently. Memory caches republish at most
568/// every [`REBALANCE_EVERY`] while they load entries, and not at all when
569/// idle (an MCP server waiting for an agent's next call).
570const BUSY_WINDOW: Duration = Duration::from_secs(5);
571
572/// How often [`auto_threads`] re-reads the registry: a burst of scans
573/// started together by several agents lasts about a second. A read checks
574/// every entry's process: about 1 ms for 31 servers on macOS, which callers
575/// wait for (letting them go on with the last count made 16 servers' cold
576/// sweep 10 % slower).
577const THREADS_EVERY: Duration = Duration::from_millis(100);
578
579/// Worker threads of a search whose caller fixed none: this machine's CPUs
580/// split between this process and the user's other busy registered ones
581/// ([`auto_threads_for`]). Re-read at most once per [`THREADS_EVERY`], so
582/// a sweep follows the others starting and finishing their scans
583/// ([`crate::Threads::wait_turn`]).
584pub fn auto_threads() -> usize {
585    static LAST: Mutex<Option<(Instant, usize)>> = Mutex::new(None);
586    let mut last = LAST.lock().unwrap();
587    if let Some((at, threads)) = *last
588        && at.elapsed() < THREADS_EVERY
589    {
590        return threads;
591    }
592    let others = match &REGISTRATION.lock().unwrap().registry {
593        Some(registry) => registry.others(),
594        // One-shot commands keep no memory cache and do not register; they
595        // still give way to the resident hosts.
596        None => ProcessRegistry::default_dir().map(|d| ProcessRegistry::others_in(&d)).unwrap_or_default(),
597    };
598    let busy = others.iter().filter(|o| o.busy.unwrap_or(o.idle < BUSY_WINDOW)).count();
599    let cpus = cpus();
600    let threads = auto_threads_for(cpus, busy);
601    *last = Some((Instant::now(), threads));
602    threads
603}
604
605/// This machine's CPUs.
606pub(crate) fn cpus() -> usize {
607    std::thread::available_parallelism().map_or(8, |n| n.get())
608}
609
610/// Workers the busy processes may start together, per CPU. An exact split
611/// (1) leaves CPUs idle while the processes holding the largest APKs finish
612/// alone: 8 servers scanning 312 APKs on 16 CPUs took 223 ms with 2 workers
613/// each, 180 ms with 8 and 176 ms with 16, for 13 % more memory at 8.
614const OVERSUBSCRIBE: usize = 4;
615
616/// [`OVERSUBSCRIBE`] x `cpus` workers split evenly between this process and
617/// `busy_others`, at most `cpus` and at least 2 (where there are 2): a scan
618/// on one thread gives up its overlap of inflating one entry while searching
619/// another. Up to 4 busy processes each keep every CPU.
620pub fn auto_threads_for(cpus: usize, busy_others: usize) -> usize {
621    (cpus * OVERSUBSCRIBE / (busy_others + 1)).clamp(cpus.min(2), cpus)
622}
623
624#[cfg(test)]
625mod tests {
626    use super::*;
627
628    #[test]
629    fn threads_split_between_busy_processes() {
630        assert_eq!(auto_threads_for(32, 0), 32);
631        assert_eq!(auto_threads_for(32, 3), 32, "4 busy processes keep every CPU");
632        assert_eq!(auto_threads_for(32, 7), 16);
633        assert_eq!(auto_threads_for(32, 15), 8);
634        assert_eq!(auto_threads_for(32, 255), 2, "never below 2");
635        assert_eq!(auto_threads_for(1, 3), 1, "nor above the CPUs");
636        assert!(auto_threads() >= 1);
637    }
638
639    #[test]
640    fn workers_wait_for_their_turn() {
641        use crate::Threads;
642        Threads::Auto.wait_turn(0, || panic!("the first worker always runs"));
643        Threads::Fixed(1).wait_turn(5, || panic!("fixed counts never hold a worker back"));
644        let asked = std::sync::atomic::AtomicUsize::new(0);
645        // Past every CPU a worker waits until the search tells it to stop.
646        Threads::Auto.wait_turn(Threads::Auto.most(), || asked.fetch_add(1, Ordering::Relaxed) >= 2);
647        assert_eq!(asked.load(Ordering::Relaxed), 3);
648    }
649
650    #[test]
651    fn results_roundtrip_and_errors_are_not_kept() {
652        let row = |rel| RefResult {
653            dex_name: "classes.dex".into(),
654            method_idx: 7,
655            method: "La;->b".into(),
656            signature: "La;->b()V".into(),
657            matched: vec![
658                Matched::String("x\ny".into()),
659                Matched::Symbol("La;->c:I".into()),
660                Matched::Field { symbol: "La;->f:I".into(), access: asc_findrefs::FieldAccess::Read },
661            ],
662            relation: rel,
663            sites: vec![
664                RefSite { offset: 9, matched: 1, field_access: None, pc: Some(3), line: None },
665                RefSite { offset: 11, matched: 2, field_access: Some(asc_findrefs::FieldAccess::Read), pc: Some(4), line: Some(8) },
666            ],
667            context: rel.map(|_| "    const-string v0, \"x\"".to_string()),
668        };
669        let batch = |name: &str, results, error: Option<String>| RefBatch {
670            entry_name: name.into(),
671            results,
672            text: None,
673            inflate_us: 5,
674            search_us: 6,
675            source: "inflate",
676            error,
677        };
678        let mut rec = ResultRecorder::new(L3_MAX_RESULT);
679        rec.record(&batch("classes.dex", vec![row(None), row(Some(ClassRelation::Implements))], None));
680        rec.record(&batch("classes2.dex", vec![], None));
681        let bytes = rec.finish(&["a note".to_string()]).unwrap();
682        let Recorded { batches: back, notes } = decode_batches(bytes.clone()).unwrap();
683        assert_eq!(notes, ["a note"]);
684        assert_eq!(back.len(), 2);
685        assert_eq!(back[0].results, vec![row(None), row(Some(ClassRelation::Implements))]);
686        assert_eq!((back[1].entry_name.as_str(), back[1].source), ("classes2.dex", "result"));
687        assert!(decode_batches(bytes[..bytes.len() - 1].to_vec()).is_none());
688        // Cut at a batch boundary: the end marker is missing.
689        let mut one = ResultRecorder::new(L3_MAX_RESULT);
690        one.record(&batch("classes2.dex", vec![], None));
691        let one = one.finish(&[]).unwrap();
692        assert!(decode_batches(one[..one.len() - 4].to_vec()).is_none());
693        assert_eq!(decode_batches(one).unwrap().batches.len(), 1);
694        assert!(decode_batches(Vec::new()).is_none());
695        let mut failed = ResultRecorder::new(L3_MAX_RESULT);
696        failed.record(&batch("classes.dex", vec![], Some("bad".into())));
697        assert!(failed.finish(&[]).is_none());
698    }
699
700    #[test]
701    fn text_rows_roundtrip_within_their_limit() {
702        let text = |lines: &[&str]| {
703            let bytes: Arc<dyn AsRef<[u8]> + Send + Sync> =
704                Arc::new(lines.iter().map(|l| format!("{l}\n")).collect::<String>().into_bytes());
705            let mut ends = Vec::new();
706            let mut at = 0;
707            for l in lines {
708                at += l.len();
709                ends.push(at);
710                at += 1;
711            }
712            let len = (*bytes).as_ref().len();
713            TextRows::shared(bytes, 0..len, ends).unwrap()
714        };
715        let batch = |name: &str, t: TextRows| RefBatch {
716            entry_name: name.into(),
717            results: Vec::new(),
718            text: Some(t),
719            inflate_us: 0,
720            search_us: 0,
721            source: "inflate",
722            error: None,
723        };
724        let a = text(&["classes.dex | La;->b()V | \"x\"", "classes.dex | La;->c()V | \"y\""]);
725        let mut rec = ResultRecorder::new(1 << 10);
726        rec.record(&batch("classes.dex", a.clone()));
727        rec.record(&batch("classes2.dex", text(&[])));
728        let Recorded { batches: back, notes } = decode_batches(rec.finish(&["n".to_string()]).unwrap()).unwrap();
729        assert_eq!(notes, ["n"]);
730        assert_eq!(back[0].text.as_ref(), Some(&a));
731        assert_eq!(back[0].row_count(), 2);
732        assert_eq!(back[1].text.as_ref().map(TextRows::len), Some(0));
733        // Past the limit nothing is kept, and the rows are not copied first.
734        let mut rec = ResultRecorder::new(64);
735        rec.record(&batch("classes.dex", a));
736        assert!(rec.out.is_empty() && rec.finish(&[]).is_none());
737    }
738}
asc_engine/tiers.rs

asc_engine/
tiers.rs
