rspack_plugin_progress/
lib.rs

1use std::{
2  cmp,
3  cmp::Ordering,
4  sync::{
5    Arc, LazyLock,
6    atomic::{AtomicU32, Ordering::Relaxed},
7  },
8  time::{Duration, Instant, SystemTime, UNIX_EPOCH},
9};
10
11use futures::future::BoxFuture;
12use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
13use rspack_collections::IdentifierMap;
14use rspack_core::{
15  BoxModule, Compilation, CompilationAfterOptimizeModules, CompilationAfterProcessAssets,
16  CompilationBuildModule, CompilationChunkIds, CompilationFinishModules, CompilationId,
17  CompilationModuleIds, CompilationOptimizeChunkModules, CompilationOptimizeChunks,
18  CompilationOptimizeCodeGeneration, CompilationOptimizeDependencies, CompilationOptimizeModules,
19  CompilationOptimizeTree, CompilationParams, CompilationProcessAssets, CompilationSeal,
20  CompilationSucceedModule, CompilerAfterEmit, CompilerClose, CompilerCompilation, CompilerEmit,
21  CompilerFinishMake, CompilerId, CompilerMake, CompilerThisCompilation, ModuleIdentifier, Plugin,
22};
23use rspack_error::Result;
24use rspack_hook::{plugin, plugin_hook};
25use tokio::sync::Mutex;
26
27type HandlerFn =
28  Arc<dyn Fn(f64, String, Vec<String>) -> BoxFuture<'static, Result<()>> + Send + Sync>;
29
30pub enum ProgressPluginOptions {
31  Handler(HandlerFn),
32  Default(ProgressPluginDisplayOptions),
33}
34
35impl std::fmt::Debug for ProgressPluginOptions {
36  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37    match self {
38      ProgressPluginOptions::Handler(_handler) => {
39        f.debug_struct("ProgressPluginOptions::Handler").finish()
40      }
41      ProgressPluginOptions::Default(options) => f
42        .debug_struct("ProgressPluginOptions::Default")
43        .field("options", &options)
44        .finish(),
45    }
46  }
47}
48
49static MULTI_PROGRESS: LazyLock<MultiProgress> =
50  LazyLock::new(|| MultiProgress::with_draw_target(ProgressDrawTarget::stdout_with_hz(100)));
51#[derive(Debug, Default)]
52pub struct ProgressPluginDisplayOptions {
53  // the prefix name of progress bar
54  pub prefix: String,
55  // tells ProgressPlugin to collect profile data for progress steps.
56  pub profile: bool,
57  // the template of progress bar, see [`indicatif::ProgressStyle::with_template`]
58  pub template: String,
59  // the tick string sequence for spinners, see [`indicatif::ProgressStyle::tick_strings`]
60  pub tick_strings: Option<Vec<String>>,
61  // the progress characters, see [`indicatif::ProgressStyle::progress_chars`]
62  pub progress_chars: String,
63}
64
65#[derive(Debug)]
66pub struct ProgressPluginStateInfo {
67  pub value: String,
68  pub time: Instant,
69  pub duration: Option<Duration>,
70}
71
72#[plugin]
73#[derive(Debug)]
74pub struct ProgressPlugin {
75  pub options: ProgressPluginOptions,
76  pub progress_bar: Option<ProgressBar>,
77  pub modules_count: Arc<AtomicU32>,
78  pub modules_done: Arc<AtomicU32>,
79  pub active_modules: Arc<Mutex<IdentifierMap<Instant>>>,
80  pub last_modules_count: Arc<Mutex<Option<u32>>>,
81  pub last_active_module: Arc<Mutex<Option<ModuleIdentifier>>>,
82  pub last_state_info: Arc<Mutex<Vec<ProgressPluginStateInfo>>>,
83  pub last_updated: Arc<AtomicU32>,
84}
85
86impl ProgressPlugin {
87  pub fn new(options: ProgressPluginOptions) -> Self {
88    let progress_bar = match &options {
89      ProgressPluginOptions::Handler(_fn) => None,
90      ProgressPluginOptions::Default(options) => {
91        // default interval is 20, means draw every 1000/20 = 50ms, use 100 to draw every 1000/100 = 10ms
92        let progress_bar = MULTI_PROGRESS.add(ProgressBar::new(100));
93
94        let mut progress_bar_style = ProgressStyle::with_template(&options.template)
95          .expect("TODO:")
96          .progress_chars(&options.progress_chars);
97        if let Some(tick_strings) = &options.tick_strings {
98          progress_bar_style = progress_bar_style.tick_strings(
99            tick_strings
100              .iter()
101              .map(|s| s.as_str())
102              .collect::<Vec<_>>()
103              .as_slice(),
104          );
105        }
106        progress_bar.set_style(progress_bar_style);
107        Some(progress_bar)
108      }
109    };
110    Self::new_inner(
111      options,
112      progress_bar,
113      Default::default(),
114      Default::default(),
115      Default::default(),
116      Default::default(),
117      Default::default(),
118      Default::default(),
119      Default::default(),
120    )
121  }
122
123  async fn update_throttled(&self) -> Result<()> {
124    let current_time = SystemTime::now()
125      .duration_since(UNIX_EPOCH)
126      .expect("failed to get current time")
127      .as_millis() as u32;
128
129    if current_time - self.last_updated.load(Relaxed) > 200 {
130      self.update().await?;
131      self.last_updated.store(current_time, Relaxed);
132    }
133
134    Ok(())
135  }
136
137  async fn update(&self) -> Result<()> {
138    let modules_done = self.modules_done.load(Relaxed);
139    let percent_by_module = (modules_done as f64)
140      / (cmp::max(
141        self.last_modules_count.lock().await.unwrap_or(1),
142        self.modules_count.load(Relaxed),
143      ) as f64);
144
145    let mut items = vec![];
146
147    if let Some(last_active_module) = self.last_active_module.lock().await.as_ref() {
148      items.push(last_active_module.to_string());
149      let duration = self
150        .active_modules
151        .lock()
152        .await
153        .get(last_active_module)
154        .map(|time| Instant::now() - *time);
155      self
156        .handler(
157          0.1 + percent_by_module * 0.55,
158          String::from("building"),
159          items,
160          duration,
161        )
162        .await?;
163    }
164    Ok(())
165  }
166
167  pub async fn handler(
168    &self,
169    percent: f64,
170    msg: String,
171    state_items: Vec<String>,
172    time: Option<Duration>,
173  ) -> Result<()> {
174    match &self.options {
175      ProgressPluginOptions::Handler(handler) => handler(percent, msg, state_items).await?,
176      ProgressPluginOptions::Default(options) => {
177        if options.profile {
178          self.default_handler(percent, msg, state_items, time).await;
179        } else {
180          self.progress_bar_handler(percent, msg, state_items);
181        }
182      }
183    };
184    Ok(())
185  }
186
187  async fn default_handler(
188    &self,
189    _: f64,
190    msg: String,
191    items: Vec<String>,
192    duration: Option<Duration>,
193  ) {
194    let full_state = [vec![msg.clone()], items.clone()].concat();
195    let now = Instant::now();
196    {
197      let mut last_state_info = self.last_state_info.lock().await;
198      let len = full_state.len().max(last_state_info.len());
199      let original_last_state_info_len = last_state_info.len();
200      for i in (0..len).rev() {
201        if i + 1 > original_last_state_info_len {
202          last_state_info.insert(
203            original_last_state_info_len,
204            ProgressPluginStateInfo {
205              value: full_state[i].clone(),
206              time: now,
207              duration: None,
208            },
209          )
210        } else if i + 1 > full_state.len() || !last_state_info[i].value.eq(&full_state[i]) {
211          let diff = match last_state_info[i].duration {
212            Some(duration) => duration,
213            _ => now - last_state_info[i].time,
214          }
215          .as_millis();
216          let report_state = if i > 0 {
217            last_state_info[i - 1].value.clone() + " > " + last_state_info[i].value.clone().as_str()
218          } else {
219            last_state_info[i].value.clone()
220          };
221
222          if diff > 5 {
223            // TODO: color map
224            let mut color = "\x1b[32m";
225            if diff > 10000 {
226              color = "\x1b[31m"
227            } else if diff > 1000 {
228              color = "\x1b[33m"
229            }
230            println!(
231              "{}{} {} ms {}\x1B[0m",
232              color,
233              " | ".repeat(i),
234              diff,
235              report_state
236            );
237          }
238          match (i + 1).cmp(&full_state.len()) {
239            Ordering::Greater => last_state_info.truncate(i),
240            Ordering::Equal => {
241              last_state_info[i] = ProgressPluginStateInfo {
242                value: full_state[i].clone(),
243                time: now,
244                duration,
245              }
246            }
247            Ordering::Less => {
248              last_state_info[i] = ProgressPluginStateInfo {
249                value: full_state[i].clone(),
250                time: now,
251                duration: None,
252              };
253            }
254          }
255        }
256      }
257    }
258  }
259
260  fn progress_bar_handler(&self, percent: f64, msg: String, state_items: Vec<String>) {
261    if let Some(progress_bar) = &self.progress_bar {
262      let msg = msg + " " + state_items.join(" ").as_str();
263      if percent == 1.0 {
264        progress_bar.finish_with_message(msg);
265      } else {
266        progress_bar.set_message(msg);
267        progress_bar.set_position((percent * 100.0) as u64);
268      }
269    }
270  }
271
272  async fn sealing_hooks_report(&self, name: &str, index: i32) -> Result<()> {
273    let number_of_sealing_hooks = 38;
274    self
275      .handler(
276        0.7 + 0.25 * (index as f64 / number_of_sealing_hooks as f64),
277        "sealing".to_string(),
278        vec![name.to_string()],
279        None,
280      )
281      .await
282  }
283
284  pub fn is_profile(&self) -> bool {
285    match &self.options {
286      ProgressPluginOptions::Handler(_) => false,
287      ProgressPluginOptions::Default(options) => options.profile,
288    }
289  }
290}
291
292#[plugin_hook(CompilerThisCompilation for ProgressPlugin)]
293async fn this_compilation(
294  &self,
295  _compilation: &mut Compilation,
296  _params: &mut CompilationParams,
297) -> Result<()> {
298  if let ProgressPluginOptions::Default(options) = &self.options {
299    let progress_bar = self.progress_bar.as_ref().unwrap_or_else(|| unreachable!());
300    if !options.profile {
301      progress_bar.reset();
302      progress_bar.set_prefix(options.prefix.clone());
303    }
304  }
305
306  self
307    .handler(
308      0.08,
309      "setup".to_string(),
310      vec!["compilation".to_string()],
311      None,
312    )
313    .await
314}
315
316#[plugin_hook(CompilerCompilation for ProgressPlugin)]
317async fn compilation(
318  &self,
319  _compilation: &mut Compilation,
320  _params: &mut CompilationParams,
321) -> Result<()> {
322  self
323    .handler(
324      0.09,
325      "setup".to_string(),
326      vec!["compilation".to_string()],
327      None,
328    )
329    .await
330}
331
332#[plugin_hook(CompilerMake for ProgressPlugin)]
333async fn make(&self, _compilation: &mut Compilation) -> Result<()> {
334  self
335    .handler(0.1, String::from("make"), vec![], None)
336    .await?;
337  self.modules_count.store(0, Relaxed);
338  self.modules_done.store(0, Relaxed);
339  Ok(())
340}
341
342#[plugin_hook(CompilationBuildModule for ProgressPlugin)]
343async fn build_module(
344  &self,
345  _compiler_id: CompilerId,
346  _compilation_id: CompilationId,
347  module: &mut BoxModule,
348) -> Result<()> {
349  self
350    .active_modules
351    .lock()
352    .await
353    .insert(module.identifier(), Instant::now());
354  self.modules_count.fetch_add(1, Relaxed);
355  self
356    .last_active_module
357    .lock()
358    .await
359    .replace(module.identifier());
360  if let ProgressPluginOptions::Default(options) = &self.options
361    && !options.profile
362  {
363    self.update_throttled().await?;
364  }
365
366  Ok(())
367}
368
369#[plugin_hook(CompilationSucceedModule for ProgressPlugin)]
370async fn succeed_module(
371  &self,
372  _compiler_id: CompilerId,
373  _compilation_id: CompilationId,
374  module: &mut BoxModule,
375) -> Result<()> {
376  self.modules_done.fetch_add(1, Relaxed);
377  self
378    .last_active_module
379    .lock()
380    .await
381    .replace(module.identifier());
382
383  // only profile mode should update at succeed module
384  if self.is_profile() {
385    self.update_throttled().await?;
386  }
387  let mut last_active_module = Default::default();
388  {
389    let mut active_modules = self.active_modules.lock().await;
390    active_modules.remove(&module.identifier());
391
392    // get the last active module
393    if !self.is_profile() {
394      active_modules.iter().for_each(|(module, _)| {
395        last_active_module = *module;
396      });
397    }
398  }
399  if !self.is_profile() {
400    self
401      .last_active_module
402      .lock()
403      .await
404      .replace(last_active_module);
405    if !last_active_module.is_empty() {
406      self.update_throttled().await?;
407    }
408  }
409  Ok(())
410}
411
412#[plugin_hook(CompilerFinishMake for ProgressPlugin)]
413async fn finish_make(&self, _compilation: &mut Compilation) -> Result<()> {
414  self
415    .handler(
416      0.69,
417      "building".to_string(),
418      vec!["finish make".to_string()],
419      None,
420    )
421    .await
422}
423
424#[plugin_hook(CompilationFinishModules for ProgressPlugin)]
425async fn finish_modules(&self, _compilation: &mut Compilation) -> Result<()> {
426  self.sealing_hooks_report("finish modules", 0).await
427}
428
429#[plugin_hook(CompilationSeal for ProgressPlugin)]
430async fn seal(&self, _compilation: &mut Compilation) -> Result<()> {
431  self.sealing_hooks_report("plugins", 1).await
432}
433
434#[plugin_hook(CompilationOptimizeDependencies for ProgressPlugin)]
435async fn optimize_dependencies(&self, _compilation: &mut Compilation) -> Result<Option<bool>> {
436  self.sealing_hooks_report("dependencies", 2).await?;
437  Ok(None)
438}
439
440#[plugin_hook(CompilationOptimizeModules for ProgressPlugin)]
441async fn optimize_modules(&self, _compilation: &mut Compilation) -> Result<Option<bool>> {
442  self.sealing_hooks_report("module optimization", 7).await?;
443  Ok(None)
444}
445
446#[plugin_hook(CompilationAfterOptimizeModules for ProgressPlugin)]
447async fn after_optimize_modules(&self, _compilation: &mut Compilation) -> Result<()> {
448  self
449    .sealing_hooks_report("after module optimization", 8)
450    .await
451}
452
453#[plugin_hook(CompilationOptimizeChunks for ProgressPlugin)]
454async fn optimize_chunks(&self, _compilation: &mut Compilation) -> Result<Option<bool>> {
455  self.sealing_hooks_report("chunk optimization", 9).await?;
456  Ok(None)
457}
458
459#[plugin_hook(CompilationOptimizeTree for ProgressPlugin)]
460async fn optimize_tree(&self, _compilation: &mut Compilation) -> Result<()> {
461  self
462    .sealing_hooks_report("module and chunk tree optimization", 11)
463    .await
464}
465
466#[plugin_hook(CompilationOptimizeChunkModules for ProgressPlugin)]
467async fn optimize_chunk_modules(&self, _compilation: &mut Compilation) -> Result<Option<bool>> {
468  self
469    .sealing_hooks_report("chunk modules optimization", 13)
470    .await?;
471  Ok(None)
472}
473
474#[plugin_hook(CompilationModuleIds for ProgressPlugin)]
475async fn module_ids(&self, _modules: &mut Compilation) -> Result<()> {
476  self.sealing_hooks_report("module ids", 16).await
477}
478
479#[plugin_hook(CompilationChunkIds for ProgressPlugin)]
480async fn chunk_ids(&self, _compilation: &mut Compilation) -> Result<()> {
481  self.sealing_hooks_report("chunk ids", 21).await
482}
483
484#[plugin_hook(CompilationOptimizeCodeGeneration for ProgressPlugin)]
485async fn optimize_code_generation(&self, _compilation: &mut Compilation) -> Result<()> {
486  self.sealing_hooks_report("code generation", 26).await
487}
488
489#[plugin_hook(CompilationProcessAssets for ProgressPlugin, stage = Compilation::PROCESS_ASSETS_STAGE_ADDITIONAL)]
490async fn process_assets(&self, _compilation: &mut Compilation) -> Result<()> {
491  self.sealing_hooks_report("asset processing", 35).await
492}
493
494#[plugin_hook(CompilationAfterProcessAssets for ProgressPlugin)]
495async fn after_process_assets(&self, _compilation: &mut Compilation) -> Result<()> {
496  self
497    .sealing_hooks_report("after asset optimization", 36)
498    .await
499}
500
501#[plugin_hook(CompilerEmit for ProgressPlugin)]
502async fn emit(&self, _compilation: &mut Compilation) -> Result<()> {
503  self
504    .handler(0.98, "emitting".to_string(), vec!["emit".to_string()], None)
505    .await
506}
507
508#[plugin_hook(CompilerAfterEmit for ProgressPlugin)]
509async fn after_emit(&self, _compilation: &mut Compilation) -> Result<()> {
510  self
511    .handler(
512      1.0,
513      "emitting".to_string(),
514      vec!["after emit".to_string()],
515      None,
516    )
517    .await
518}
519
520#[plugin_hook(CompilerClose for ProgressPlugin)]
521async fn close(&self, _compilation: &Compilation) -> Result<()> {
522  if let Some(progress_bar) = &self.progress_bar {
523    MULTI_PROGRESS.remove(progress_bar);
524  }
525  Ok(())
526}
527
528impl Plugin for ProgressPlugin {
529  fn name(&self) -> &'static str {
530    "progress"
531  }
532
533  fn apply(&self, ctx: &mut rspack_core::ApplyContext<'_>) -> Result<()> {
534    ctx
535      .compiler_hooks
536      .this_compilation
537      .tap(this_compilation::new(self));
538    ctx.compiler_hooks.compilation.tap(compilation::new(self));
539    ctx.compiler_hooks.make.tap(make::new(self));
540    ctx
541      .compilation_hooks
542      .build_module
543      .tap(build_module::new(self));
544    ctx
545      .compilation_hooks
546      .succeed_module
547      .tap(succeed_module::new(self));
548    ctx.compiler_hooks.finish_make.tap(finish_make::new(self));
549    ctx
550      .compilation_hooks
551      .finish_modules
552      .tap(finish_modules::new(self));
553    ctx.compilation_hooks.seal.tap(seal::new(self));
554    ctx
555      .compilation_hooks
556      .optimize_dependencies
557      .tap(optimize_dependencies::new(self));
558    ctx
559      .compilation_hooks
560      .optimize_modules
561      .tap(optimize_modules::new(self));
562    ctx
563      .compilation_hooks
564      .after_optimize_modules
565      .tap(after_optimize_modules::new(self));
566    ctx
567      .compilation_hooks
568      .optimize_chunks
569      .tap(optimize_chunks::new(self));
570    ctx
571      .compilation_hooks
572      .optimize_tree
573      .tap(optimize_tree::new(self));
574    ctx
575      .compilation_hooks
576      .optimize_chunk_modules
577      .tap(optimize_chunk_modules::new(self));
578    ctx.compilation_hooks.module_ids.tap(module_ids::new(self));
579    ctx.compilation_hooks.chunk_ids.tap(chunk_ids::new(self));
580    ctx
581      .compilation_hooks
582      .optimize_code_generation
583      .tap(optimize_code_generation::new(self));
584    ctx
585      .compilation_hooks
586      .process_assets
587      .tap(process_assets::new(self));
588    ctx
589      .compilation_hooks
590      .after_process_assets
591      .tap(after_process_assets::new(self));
592    ctx.compiler_hooks.emit.tap(emit::new(self));
593    ctx.compiler_hooks.after_emit.tap(after_emit::new(self));
594    ctx.compiler_hooks.close.tap(close::new(self));
595    Ok(())
596  }
597}