ide_db/
prime_caches.rs

1//! rust-analyzer is lazy and doesn't compute anything unless asked. This
2//! sometimes is counter productive when, for example, the first goto definition
3//! request takes longer to compute. This module implements prepopulation of
4//! various caches, it's not really advanced at the moment.
5use std::panic::AssertUnwindSafe;
6
7use hir::{Symbol, import_map::ImportMap, sym};
8use rustc_hash::{FxHashMap, FxHashSet};
9use salsa::{Cancelled, Database};
10
11use crate::{FxIndexMap, RootDatabase, base_db::Crate, symbol_index::SymbolIndex};
12
13/// We're indexing many crates.
14#[derive(Debug)]
15pub struct ParallelPrimeCachesProgress {
16    /// the crates that we are currently priming.
17    pub crates_currently_indexing: Vec<Symbol>,
18    /// the total number of crates we want to prime.
19    pub crates_total: usize,
20    /// the total number of crates that have finished priming
21    pub crates_done: usize,
22    pub work_type: &'static str,
23}
24
25/// Warm caches for `scope`.
26///
27/// `scope` must be closed under transitive dependencies: the scheduler only
28/// follows reverse-dep edges within `scope`. Out-of-scope crates can still be
29/// primed by salsa on demand when a scope crate's queries reach into them.
30/// Callers that want to prime everything pass `&all_crates(db)`.
31pub fn parallel_prime_caches(
32    db: &RootDatabase,
33    scope: &[Crate],
34    num_worker_threads: usize,
35    cb: &(dyn Fn(ParallelPrimeCachesProgress) + Sync),
36) {
37    if scope.is_empty() {
38        return;
39    }
40    let _p = tracing::info_span!("parallel_prime_caches", scope_size = scope.len()).entered();
41
42    enum ParallelPrimeCacheWorkerProgress {
43        BeginCrateDefMap { crate_id: Crate, crate_name: Symbol },
44        EndCrateDefMap { crate_id: Crate },
45        EndCrateImportMap,
46        EndSema,
47        EndModuleSymbols,
48        Cancelled(Cancelled),
49    }
50
51    // The setup here is a bit complicated. We try to make best use of compute resources.
52    // The idea is that if we have a def map available to compute, we should do that first.
53    // This is because def map is a dependency of both import map and symbols. So if we have
54    // e.g. a def map and a symbols, if we compute the def map we can, after it completes,
55    // compute the def maps of dependencies, the existing symbols and the symbols of the
56    // new crate, all in parallel. But if we compute the symbols, after that we will only
57    // have the def map to compute, and the rest of the CPU cores will rest, which is not
58    // good.
59    // However, it's better to compute symbols/import map than to compute a def map that
60    // isn't ready yet, because one of its dependencies hasn't yet completed its def map.
61    // Such def map will just block on the dependency, which is just wasted time. So better
62    // to compute the symbols/import map of an already computed def map in that time.
63
64    let scope_set: FxHashSet<Crate> = scope.iter().copied().collect();
65
66    let (reverse_deps, mut to_be_done_deps) = {
67        // Only count in-scope deps — otherwise an out-of-scope dep would
68        // leave the scheduler waiting on a crate it never enqueued.
69        let to_be_done_deps = scope
70            .iter()
71            .map(|&krate| {
72                let count = krate
73                    .data(db)
74                    .dependencies
75                    .iter()
76                    .filter(|dep| scope_set.contains(&dep.crate_id))
77                    .count() as u32;
78                (krate, count)
79            })
80            .collect::<FxHashMap<_, _>>();
81        let mut reverse_deps =
82            scope.iter().map(|&krate| (krate, Vec::new())).collect::<FxHashMap<_, _>>();
83        for &krate in scope {
84            for dep in &krate.data(db).dependencies {
85                if let Some(rev) = reverse_deps.get_mut(&dep.crate_id) {
86                    rev.push(krate);
87                }
88            }
89        }
90        (reverse_deps, to_be_done_deps)
91    };
92
93    let (
94        def_map_work_sender,
95        import_map_work_sender,
96        symbols_work_sender,
97        sema_work_sender,
98        progress_receiver,
99    ) = {
100        let (progress_sender, progress_receiver) = crossbeam_channel::unbounded();
101        let (def_map_work_sender, def_map_work_receiver) = crossbeam_channel::unbounded();
102        let (import_map_work_sender, import_map_work_receiver) = crossbeam_channel::unbounded();
103        let (sema_work_sender, sema_work_receiver) = crossbeam_channel::unbounded();
104        let (symbols_work_sender, symbols_work_receiver) = crossbeam_channel::unbounded();
105        let prime_caches_worker = move |db: RootDatabase| {
106            let handle_def_map = |crate_id, crate_name| {
107                progress_sender.send(ParallelPrimeCacheWorkerProgress::BeginCrateDefMap {
108                    crate_id,
109                    crate_name,
110                })?;
111
112                let cancelled = Cancelled::catch(|| {
113                    _ = hir::crate_def_map(&db, crate_id);
114                });
115
116                match cancelled {
117                    Ok(()) => progress_sender
118                        .send(ParallelPrimeCacheWorkerProgress::EndCrateDefMap { crate_id })?,
119                    Err(cancelled) => progress_sender
120                        .send(ParallelPrimeCacheWorkerProgress::Cancelled(cancelled))?,
121                }
122
123                Ok::<_, crossbeam_channel::SendError<_>>(())
124            };
125            let handle_sema = |crate_id| {
126                let cancelled = Cancelled::catch(|| {
127                    hir::attach_db(&db, || {
128                        // method resolution is likely to hit all trait impls at some point
129                        // we pre-populate it here as this will hit a lot of parses ...
130                        // This also computes the lang items, which is what we want as the work for them is also highly recursive and will be trigger by the module symbols query
131                        // slowing down leaf crate analysis tremendously as we go back to being blocked on a single thread
132                        _ = hir::TraitImpls::for_crate(&db, crate_id);
133                    })
134                });
135
136                match cancelled {
137                    Ok(()) => progress_sender.send(ParallelPrimeCacheWorkerProgress::EndSema)?,
138                    Err(cancelled) => progress_sender
139                        .send(ParallelPrimeCacheWorkerProgress::Cancelled(cancelled))?,
140                }
141
142                Ok::<_, crossbeam_channel::SendError<_>>(())
143            };
144            let handle_import_map = |crate_id| {
145                let cancelled = Cancelled::catch(|| _ = ImportMap::of(&db, crate_id));
146
147                match cancelled {
148                    Ok(()) => {
149                        progress_sender.send(ParallelPrimeCacheWorkerProgress::EndCrateImportMap)?
150                    }
151                    Err(cancelled) => progress_sender
152                        .send(ParallelPrimeCacheWorkerProgress::Cancelled(cancelled))?,
153                }
154
155                Ok::<_, crossbeam_channel::SendError<_>>(())
156            };
157            let handle_symbols = |module: hir::Module| {
158                let cancelled = Cancelled::catch(AssertUnwindSafe(|| {
159                    _ = SymbolIndex::module_symbols(&db, module)
160                }));
161
162                match cancelled {
163                    Ok(()) => {
164                        progress_sender.send(ParallelPrimeCacheWorkerProgress::EndModuleSymbols)?
165                    }
166                    Err(cancelled) => progress_sender
167                        .send(ParallelPrimeCacheWorkerProgress::Cancelled(cancelled))?,
168                }
169
170                Ok::<_, crossbeam_channel::SendError<_>>(())
171            };
172
173            loop {
174                db.unwind_if_revision_cancelled();
175
176                // Biased because we want to prefer def maps.
177                crossbeam_channel::select_biased! {
178                    recv(def_map_work_receiver) -> work => {
179                        let Ok((crate_id, crate_name)) = work else { break };
180                        handle_def_map(crate_id, crate_name)?;
181                    }
182                    recv(sema_work_receiver) -> work => {
183                        let Ok(crate_id) = work else { break };
184                        handle_sema(crate_id)?;
185                    }
186                    recv(import_map_work_receiver) -> work => {
187                        let Ok(crate_id) = work else { break };
188                        handle_import_map(crate_id)?;
189                    }
190                    recv(symbols_work_receiver) -> work => {
191                        let Ok(module) = work else { break };
192                        handle_symbols(module)?;
193                    }
194                }
195            }
196            Ok::<_, crossbeam_channel::SendError<_>>(())
197        };
198
199        for id in 0..num_worker_threads {
200            stdx::thread::Builder::new(
201                stdx::thread::ThreadIntent::Worker,
202                format!("PrimeCaches#{id}"),
203            )
204            .allow_leak(true)
205            .spawn({
206                let worker = prime_caches_worker.clone();
207                let db = db.clone();
208                move || worker(db)
209            })
210            .expect("failed to spawn thread");
211        }
212
213        (
214            def_map_work_sender,
215            import_map_work_sender,
216            symbols_work_sender,
217            sema_work_sender,
218            progress_receiver,
219        )
220    };
221
222    let crate_def_maps_total = scope.len();
223    let mut crate_def_maps_done = 0;
224    let (mut crate_import_maps_total, mut crate_import_maps_done) = (0usize, 0usize);
225    let (mut module_symbols_total, mut module_symbols_done) = (0usize, 0usize);
226    let (mut sema_total, mut sema_done) = (0usize, 0usize);
227
228    // an index map is used to preserve ordering so we can sort the progress report in order of
229    // "longest crate to index" first
230    let mut crates_currently_indexing =
231        FxIndexMap::with_capacity_and_hasher(num_worker_threads, Default::default());
232
233    for (&krate, &to_be_done_deps) in &to_be_done_deps {
234        if to_be_done_deps != 0 {
235            continue;
236        }
237
238        let name = crate_name(db, krate);
239        def_map_work_sender.send((krate, name)).ok();
240    }
241
242    while crate_def_maps_done < crate_def_maps_total
243        || crate_import_maps_done < crate_import_maps_total
244        || module_symbols_done < module_symbols_total
245        || sema_done < sema_total
246    {
247        db.unwind_if_revision_cancelled();
248
249        let progress = ParallelPrimeCachesProgress {
250            crates_currently_indexing: crates_currently_indexing.values().cloned().collect(),
251            crates_done: crate_def_maps_done,
252            crates_total: crate_def_maps_total,
253            work_type: "Indexing",
254        };
255
256        cb(progress);
257
258        // Biased to prefer progress updates (and because it's faster).
259        let progress = match progress_receiver.recv() {
260            Ok(p) => p,
261            Err(crossbeam_channel::RecvError) => {
262                // all our workers have exited, mark us as finished and exit
263                cb(ParallelPrimeCachesProgress {
264                    crates_currently_indexing: vec![],
265                    crates_done: crate_def_maps_done,
266                    crates_total: crate_def_maps_done,
267                    work_type: "Done",
268                });
269                return;
270            }
271        };
272
273        match progress {
274            ParallelPrimeCacheWorkerProgress::BeginCrateDefMap { crate_id, crate_name } => {
275                crates_currently_indexing.insert(crate_id, crate_name);
276            }
277            ParallelPrimeCacheWorkerProgress::EndCrateDefMap { crate_id } => {
278                crates_currently_indexing.swap_remove(&crate_id);
279                crate_def_maps_done += 1;
280
281                // Fire ready dependencies.
282                for &dep in &reverse_deps[&crate_id] {
283                    let to_be_done = to_be_done_deps.get_mut(&dep).unwrap();
284                    *to_be_done -= 1;
285                    if *to_be_done == 0 {
286                        let dep_name = crate_name(db, dep);
287                        def_map_work_sender.send((dep, dep_name)).ok();
288                    }
289                }
290
291                if crate_def_maps_done == crate_def_maps_total {
292                    cb(ParallelPrimeCachesProgress {
293                        crates_currently_indexing: vec![],
294                        crates_done: crate_def_maps_done,
295                        crates_total: crate_def_maps_done,
296                        work_type: "Collecting Symbols",
297                    });
298                }
299
300                sema_work_sender.send(crate_id).ok();
301                sema_total += 1;
302                let origin = &crate_id.data(db).origin;
303                if origin.is_lang() {
304                    crate_import_maps_total += 1;
305                    import_map_work_sender.send(crate_id).ok();
306                } else if origin.is_local() {
307                    // Compute the symbol search index.
308                    // This primes the cache for `ide_db::symbol_index::world_symbols()`.
309                    //
310                    // We do this for workspace crates only (members of local_roots), because doing it
311                    // for all dependencies could be *very* unnecessarily slow in a large project.
312                    //
313                    // FIXME: We should do it unconditionally if the configuration is set to default to
314                    // searching dependencies (rust-analyzer.workspace.symbol.search.scope), but we
315                    // would need to pipe that configuration information down here.
316                    let modules = hir::Crate::from(crate_id).modules(db);
317                    module_symbols_total += modules.len();
318                    for module in modules {
319                        symbols_work_sender.send(module).ok();
320                    }
321                }
322            }
323            ParallelPrimeCacheWorkerProgress::EndCrateImportMap => crate_import_maps_done += 1,
324            ParallelPrimeCacheWorkerProgress::EndModuleSymbols => module_symbols_done += 1,
325            ParallelPrimeCacheWorkerProgress::EndSema => sema_done += 1,
326            ParallelPrimeCacheWorkerProgress::Cancelled(cancelled) => {
327                // Cancelled::throw should probably be public
328                std::panic::resume_unwind(Box::new(cancelled));
329            }
330        }
331    }
332}
333
334fn crate_name(db: &RootDatabase, krate: Crate) -> Symbol {
335    krate
336        .extra_data(db)
337        .display_name
338        .as_deref()
339        .cloned()
340        .unwrap_or_else(|| sym::Integer::get(salsa::plumbing::AsId::as_id(&krate).index() as usize))
341}
ide_db/prime_caches.rs

ide_db/
prime_caches.rs
