Skip to content

Commit 6fa168a

Browse files
ChayimFriedman2ChayimFriedman2
committed
Add a --num-threads to the prime-caches CLI command
And make it parallel by default (and remove the `--parallel` flag) to mirror the IDE cache priming.
1 parent 2636676 commit 6fa168a

File tree

3 files changed

+61
-60
lines changed

3 files changed

+61
-60
lines changed

‎src/tools/rust-analyzer/crates/ide-db/src/prime_caches.rs‎

Lines changed: 57 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -56,14 +56,15 @@ pub fn parallel_prime_caches(
5656

5757
let (reverse_deps, mut to_be_done_deps) = {
5858
let all_crates = db.all_crates();
59+
let to_be_done_deps = all_crates
60+
.iter()
61+
.map(|&krate| (krate, krate.data(db).dependencies.len() as u32))
62+
.collect::<FxHashMap<_, _>>();
5963
let mut reverse_deps =
6064
all_crates.iter().map(|&krate| (krate, Vec::new())).collect::<FxHashMap<_, _>>();
61-
let mut to_be_done_deps =
62-
all_crates.iter().map(|&krate| (krate, 0u32)).collect::<FxHashMap<_, _>>();
6365
for &krate in &*all_crates {
6466
for dep in &krate.data(db).dependencies {
6567
reverse_deps.get_mut(&dep.crate_id).unwrap().push(krate);
66-
*to_be_done_deps.get_mut(&krate).unwrap() += 1;
6768
}
6869
}
6970
(reverse_deps, to_be_done_deps)
@@ -74,72 +75,72 @@ pub fn parallel_prime_caches(
7475
let (def_map_work_sender, def_map_work_receiver) = crossbeam_channel::unbounded();
7576
let (import_map_work_sender, import_map_work_receiver) = crossbeam_channel::unbounded();
7677
let (symbols_work_sender, symbols_work_receiver) = crossbeam_channel::unbounded();
77-
let prime_caches_worker = move |db: RootDatabase| {
78-
let handle_def_map = |crate_id, crate_name| {
79-
progress_sender.send(ParallelPrimeCacheWorkerProgress::BeginCrateDefMap {
80-
crate_id,
81-
crate_name,
82-
})?;
78+
let prime_caches_worker =
79+
move |db: RootDatabase| {
80+
let handle_def_map = |crate_id, crate_name| {
81+
progress_sender.send(ParallelPrimeCacheWorkerProgress::BeginCrateDefMap {
82+
crate_id,
83+
crate_name,
84+
})?;
8385

84-
let cancelled = Cancelled::catch(|| _ = hir::crate_def_map(&db, crate_id));
86+
let cancelled = Cancelled::catch(|| _ = hir::crate_def_map(&db, crate_id));
8587

86-
match cancelled {
87-
Ok(()) => progress_sender
88-
.send(ParallelPrimeCacheWorkerProgress::EndCrateDefMap { crate_id })?,
89-
Err(cancelled) => progress_sender
90-
.send(ParallelPrimeCacheWorkerProgress::Cancelled(cancelled))?,
91-
}
88+
match cancelled {
89+
Ok(()) => progress_sender
90+
.send(ParallelPrimeCacheWorkerProgress::EndCrateDefMap { crate_id })?,
91+
Err(cancelled) => progress_sender
92+
.send(ParallelPrimeCacheWorkerProgress::Cancelled(cancelled))?,
93+
}
9294

93-
Ok::<_, crossbeam_channel::SendError<_>>(())
94-
};
95-
let handle_import_map = |crate_id| {
96-
let cancelled = Cancelled::catch(|| _ = db.import_map(crate_id));
95+
Ok::<_, crossbeam_channel::SendError<_>>(())
96+
};
97+
let handle_import_map = |crate_id| {
98+
let cancelled = Cancelled::catch(|| _ = db.import_map(crate_id));
9799

98-
match cancelled {
99-
Ok(()) => {
100-
progress_sender.send(ParallelPrimeCacheWorkerProgress::EndCrateImportMap)?
100+
match cancelled {
101+
Ok(()) => progress_sender
102+
.send(ParallelPrimeCacheWorkerProgress::EndCrateImportMap)?,
103+
Err(cancelled) => progress_sender
104+
.send(ParallelPrimeCacheWorkerProgress::Cancelled(cancelled))?,
101105
}
102-
Err(cancelled) => progress_sender
103-
.send(ParallelPrimeCacheWorkerProgress::Cancelled(cancelled))?,
104-
}
105106

106-
Ok::<_, crossbeam_channel::SendError<_>>(())
107-
};
108-
let handle_symbols = |module| {
109-
let cancelled =
110-
Cancelled::catch(AssertUnwindSafe(|| _ = db.module_symbols(module)));
107+
Ok::<_, crossbeam_channel::SendError<_>>(())
108+
};
109+
let handle_symbols = |module| {
110+
let cancelled =
111+
Cancelled::catch(AssertUnwindSafe(|| _ = db.module_symbols(module)));
111112

112-
match cancelled {
113-
Ok(()) => {
114-
progress_sender.send(ParallelPrimeCacheWorkerProgress::EndModuleSymbols)?
113+
match cancelled {
114+
Ok(()) => progress_sender
115+
.send(ParallelPrimeCacheWorkerProgress::EndModuleSymbols)?,
116+
Err(cancelled) => progress_sender
117+
.send(ParallelPrimeCacheWorkerProgress::Cancelled(cancelled))?,
115118
}
116-
Err(cancelled) => progress_sender
117-
.send(ParallelPrimeCacheWorkerProgress::Cancelled(cancelled))?,
118-
}
119119

120-
Ok::<_, crossbeam_channel::SendError<_>>(())
121-
};
120+
Ok::<_, crossbeam_channel::SendError<_>>(())
121+
};
122122

123-
loop {
124-
db.unwind_if_revision_cancelled();
123+
loop {
124+
db.unwind_if_revision_cancelled();
125125

126-
// Biased because we want to prefer def maps.
127-
crossbeam_channel::select_biased! {
128-
recv(def_map_work_receiver) -> work => {
129-
let Ok((crate_id, crate_name)) = work else { return Ok::<_, crossbeam_channel::SendError<_>>(()) };
130-
handle_def_map(crate_id, crate_name)?;
131-
}
132-
recv(import_map_work_receiver) -> work => {
133-
let Ok(crate_id) = work else { return Ok(()) };
134-
handle_import_map(crate_id)?;
135-
}
136-
recv(symbols_work_receiver) -> work => {
137-
let Ok(module) = work else { return Ok(()) };
138-
handle_symbols(module)?;
126+
// Biased because we want to prefer def maps.
127+
crossbeam_channel::select_biased! {
128+
recv(def_map_work_receiver) -> work => {
129+
let Ok((crate_id, crate_name)) = work else { break };
130+
handle_def_map(crate_id, crate_name)?;
131+
}
132+
recv(import_map_work_receiver) -> work => {
133+
let Ok(crate_id) = work else { break };
134+
handle_import_map(crate_id)?;
135+
}
136+
recv(symbols_work_receiver) -> work => {
137+
let Ok(module) = work else { break };
138+
handle_symbols(module)?;
139+
}
139140
}
140141
}
141-
}
142-
};
142+
Ok::<_, crossbeam_channel::SendError<_>>(())
143+
};
143144

144145
for id in 0..num_worker_threads {
145146
stdx::thread::Builder::new(

‎src/tools/rust-analyzer/crates/rust-analyzer/src/cli/flags.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -150,8 +150,8 @@ xflags::xflags! {
150150
optional --disable-proc-macros
151151
/// Run the proc-macro-srv binary at the specified path.
152152
optional --proc-macro-srv path: PathBuf
153-
/// Run cache priming in parallel.
154-
optional --parallel
153+
/// The number of threads to use. Defaults to the number of physical cores.
154+
optional --num-threads num_threads: usize
155155
}
156156

157157
cmd ssr {
@@ -299,7 +299,7 @@ pub struct PrimeCaches {
299299
pub disable_build_scripts: bool,
300300
pub disable_proc_macros: bool,
301301
pub proc_macro_srv: Option<PathBuf>,
302-
pub parallel: bool,
302+
pub num_threads: Option<usize>,
303303
}
304304

305305
#[derive(Debug)]

‎src/tools/rust-analyzer/crates/rust-analyzer/src/cli/prime_caches.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ impl flags::PrimeCaches {
5252
elapsed.memory.allocated.megabytes() as u64
5353
);
5454

55-
let threads = if self.parallel { num_cpus::get() } else { 1 };
55+
let threads = self.num_threads.unwrap_or_else(num_cpus::get_physical);
5656
ide_db::prime_caches::parallel_prime_caches(&db, threads, &|_| ());
5757

5858
let elapsed = stop_watch.elapsed();

0 commit comments

Comments
 (0)