feat: improve rs plugin and js hooks (#855)

* feat: more js hooks * docs: update docs/config.md * chore: code style * chore: code style * docs: update * chore: on_generate_file > _on_generate_file * chore: code style
umijs · Jan 11, 2024 · c1dbfbf · c1dbfbf
1 parent 1e9dcdf
commit c1dbfbf
Show file tree

Hide file tree

Showing 17 changed files with 508 additions and 335 deletions.
diff --git a/crates/mako/src/build.rs b/crates/mako/src/build.rs
@@ -53,6 +53,7 @@ impl Compiler {
     pub fn build(&self) -> Result<()> {
         debug!("build");
         let t_build = Instant::now();
+        self.context.plugin_driver.build_start(&self.context)?;
         let entries: Vec<&PathBuf> = self.context.config.entry.values().collect();
         let tasks = entries
             .iter()

diff --git a/crates/mako/src/compiler.rs b/crates/mako/src/compiler.rs
@@ -2,7 +2,7 @@ use std::collections::HashMap;
 use std::fs;
 use std::path::PathBuf;
 use std::sync::{Arc, Mutex, RwLock};
-use std::time::Instant;
+use std::time::{Instant, UNIX_EPOCH};
 
 use mako_core::anyhow::{anyhow, Error, Result};
 use mako_core::colored::Colorize;
@@ -16,7 +16,7 @@ use crate::comments::Comments;
 use crate::config::{hash_config, Config, OutputMode};
 use crate::module_graph::ModuleGraph;
 use crate::optimize_chunk::OptimizeChunksInfo;
-use crate::plugin::{Plugin, PluginDriver};
+use crate::plugin::{Plugin, PluginDriver, PluginGenerateEndParams, PluginGenerateStats};
 use crate::plugins;
 use crate::plugins::minifish::Inject;
 use crate::resolve::{get_resolvers, Resolvers};
@@ -361,6 +361,7 @@ impl Compiler {
         }
 
         let t_compiler = Instant::now();
+        let start_time = std::time::SystemTime::now();
         let is_prod = self.context.config.mode == crate::config::Mode::Production;
         let building_with_message = format!(
             "Building with {} for {}...",
@@ -377,19 +378,32 @@ impl Compiler {
             mako_core::mako_profile_scope!("Generate Stage");
             self.generate()
         };
-        let t_compiler = t_compiler.elapsed();
+        let t_compiler_duration = t_compiler.elapsed();
         if result.is_ok() {
             println!(
                 "{}",
                 format!(
                     "✓ Built in {}",
-                    format!("{}ms", t_compiler.as_millis()).bold()
+                    format!("{}ms", t_compiler_duration.as_millis()).bold()
                 )
                 .green()
             );
             if !self.context.args.watch {
                 println!("{}", "Complete!".bold());
             }
+            let end_time = std::time::SystemTime::now();
+            let params = PluginGenerateEndParams {
+                is_first_compile: true,
+                time: t_compiler.elapsed().as_millis() as u64,
+                stats: PluginGenerateStats {
+                    start_time: start_time.duration_since(UNIX_EPOCH).unwrap().as_millis() as u64,
+                    end_time: end_time.duration_since(UNIX_EPOCH).unwrap().as_millis() as u64,
+                },
+            };
+            self.context
+                .plugin_driver
+                .generate_end(&params, &self.context)
+                .unwrap();
             Ok(())
         } else {
             result

diff --git a/crates/mako/src/dev.rs b/crates/mako/src/dev.rs
@@ -15,6 +15,7 @@ use mako_core::tungstenite::Message;
 use mako_core::{hyper, hyper_staticfile, hyper_tungstenite, tokio};
 
 use crate::compiler::{Compiler, Context};
+use crate::plugin::{PluginGenerateEndParams, PluginGenerateStats};
 use crate::watch::Watch;
 
 pub struct DevServer {
@@ -281,6 +282,20 @@ impl DevServer {
             );
 
             let end_time = std::time::SystemTime::now();
+            let params = PluginGenerateEndParams {
+                is_first_compile: false,
+                time: t_compiler.elapsed().as_millis() as u64,
+                stats: PluginGenerateStats {
+                    start_time: start_time.duration_since(UNIX_EPOCH).unwrap().as_millis() as u64,
+                    end_time: end_time.duration_since(UNIX_EPOCH).unwrap().as_millis() as u64,
+                },
+            };
+            compiler
+                .context
+                .plugin_driver
+                .generate_end(&params, &compiler.context)
+                .unwrap();
+            // TODO: remove this?
             callback(OnDevCompleteParams {
                 is_first_compile: false,
                 time: t_compiler.elapsed().as_millis() as u64,

diff --git a/crates/mako/src/plugin.rs b/crates/mako/src/plugin.rs
@@ -40,6 +40,19 @@ pub struct PluginDepAnalyzeParam<'a> {
     pub ast: &'a ModuleAst,
 }
 
+#[derive(Clone)]
+pub struct PluginGenerateEndParams {
+    pub is_first_compile: bool,
+    pub time: u64,
+    pub stats: PluginGenerateStats,
+}
+
+#[derive(Clone)]
+pub struct PluginGenerateStats {
+    pub start_time: u64,
+    pub end_time: u64,
+}
+
 pub trait Plugin: Any + Send + Sync {
     fn name(&self) -> &str;
 
@@ -101,6 +114,18 @@ pub trait Plugin: Any + Send + Sync {
         Ok(None)
     }
 
+    fn build_start(&self, _context: &Arc<Context>) -> Result<Option<()>> {
+        Ok(None)
+    }
+
+    fn generate_end(
+        &self,
+        _params: &PluginGenerateEndParams,
+        _context: &Arc<Context>,
+    ) -> Result<Option<()>> {
+        Ok(None)
+    }
+
     fn runtime_plugins(&self, _context: &Arc<Context>) -> Result<Vec<String>> {
         Ok(Vec::new())
     }
@@ -224,6 +249,24 @@ impl PluginDriver {
         Err(anyhow!("None of the plugins generate content"))
     }
 
+    pub fn build_start(&self, context: &Arc<Context>) -> Result<Option<()>> {
+        for plugin in &self.plugins {
+            plugin.build_start(context)?;
+        }
+        Ok(None)
+    }
+
+    pub fn generate_end(
+        &self,
+        param: &PluginGenerateEndParams,
+        context: &Arc<Context>,
+    ) -> Result<Option<()>> {
+        for plugin in &self.plugins {
+            plugin.generate_end(param, context)?;
+        }
+        Ok(None)
+    }
+
     pub fn build_success(
         &self,
         stats: &StatsJsonMap,

diff --git a/crates/node/src/js_plugin.rs b/crates/node/src/js_plugin.rs
@@ -0,0 +1,92 @@
+use std::sync::{mpsc, Arc};
+
+use crate::threadsafe_function;
+use crate::tsfn::{LoadResult, ReadMessage, TsFnHooks, WriteRequest};
+
+pub struct JsPlugin {
+    pub hooks: TsFnHooks,
+}
+use mako::compiler::Context;
+use mako::load::Content;
+use mako::plugin::{Plugin, PluginGenerateEndParams, PluginLoadParam};
+use mako_core::anyhow::{anyhow, Result};
+
+impl Plugin for JsPlugin {
+    fn name(&self) -> &str {
+        "js_plugin"
+    }
+
+    fn build_start(&self, _context: &Arc<Context>) -> Result<Option<()>> {
+        if let Some(hook) = &self.hooks.build_start {
+            let (tx, rx) = mpsc::channel::<napi::Result<()>>();
+            hook.call(
+                ReadMessage { message: (), tx },
+                threadsafe_function::ThreadsafeFunctionCallMode::Blocking,
+            );
+            rx.recv()
+                .unwrap_or_else(|e| panic!("recv error: {:?}", e.to_string()))?;
+        }
+        Ok(None)
+    }
+
+    fn load(&self, param: &PluginLoadParam, _context: &Arc<Context>) -> Result<Option<Content>> {
+        if let Some(hook) = &self.hooks.load {
+            let (tx, rx) = mpsc::channel::<napi::Result<Option<LoadResult>>>();
+            hook.call(
+                ReadMessage {
+                    message: param.task.request.path.clone(),
+                    tx,
+                },
+                threadsafe_function::ThreadsafeFunctionCallMode::Blocking,
+            );
+            let x = rx
+                .recv()
+                .unwrap_or_else(|e| panic!("recv error: {:?}", e.to_string()))?;
+            if let Some(x) = x {
+                match x.content_type.as_str() {
+                    "js" => return Ok(Some(Content::Js(x.content))),
+                    "css" => return Ok(Some(Content::Css(x.content))),
+                    _ => return Err(anyhow!("Unsupported content type: {}", x.content_type)),
+                }
+            }
+        }
+        Ok(None)
+    }
+
+    fn generate_end(
+        &self,
+        param: &PluginGenerateEndParams,
+        _context: &Arc<Context>,
+    ) -> Result<Option<()>> {
+        if let Some(hook) = &self.hooks.generate_end {
+            let (tx, rx) = mpsc::channel::<napi::Result<()>>();
+            hook.call(
+                ReadMessage {
+                    message: param.clone(),
+                    tx,
+                },
+                threadsafe_function::ThreadsafeFunctionCallMode::Blocking,
+            );
+            rx.recv()
+                .unwrap_or_else(|e| panic!("recv error: {:?}", e.to_string()))?;
+        }
+        Ok(None)
+    }
+
+    fn before_write_fs(&self, path: &std::path::Path, content: &[u8]) -> Result<()> {
+        if let Some(hook) = &self.hooks._on_generate_file {
+            let (tx, rx) = mpsc::channel::<napi::Result<()>>();
+            hook.call(
+                WriteRequest {
+                    path: path.to_path_buf(),
+                    content: content.to_vec(),
+                    tx,
+                },
+                threadsafe_function::ThreadsafeFunctionCallMode::Blocking,
+            );
+            rx.recv()
+                .unwrap_or_else(|e| panic!("recv error: {:?}", e.to_string()))?;
+        }
+        Ok(())
+    }
+}