Skip to content

Commit

Permalink
feat(compilation): add queue handlers representing task_queue
Browse files Browse the repository at this point in the history
  • Loading branch information
JSerFeng committed Nov 27, 2023
1 parent 8b0915d commit 991cb90
Show file tree
Hide file tree
Showing 14 changed files with 169 additions and 18 deletions.
2 changes: 2 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

7 changes: 5 additions & 2 deletions crates/node_binding/src/plugins/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,9 @@ use rspack_binding_values::JsExecuteModuleArg;
use rspack_binding_values::{AfterResolveData, JsChunkAssetArgs, JsModule};
use rspack_binding_values::{BeforeResolveData, JsAssetEmittedArgs, ToJsModule};
use rspack_binding_values::{JsResolveForSchemeInput, JsResolveForSchemeResult};
use rspack_core::{ChunkAssetArgs, ModuleIdentifier, NormalModuleAfterResolveArgs};
use rspack_core::{
ChunkAssetArgs, ModuleIdentifier, NormalModuleAfterResolveArgs, SucceedModuleArgs,
};
use rspack_core::{NormalModuleBeforeResolveArgs, PluginNormalModuleFactoryAfterResolveOutput};
use rspack_core::{PluginNormalModuleFactoryBeforeResolveOutput, ResourceData};
use rspack_core::{PluginNormalModuleFactoryResolveForSchemeOutput, PluginShouldEmitHookOutput};
Expand Down Expand Up @@ -733,11 +735,12 @@ impl rspack_core::Plugin for JsHooksAdapter {
.map_err(|err| internal_error!("Failed to call after emit: {err}"))?
}

async fn succeed_module(&self, args: &dyn rspack_core::Module) -> rspack_error::Result<()> {
async fn succeed_module(&self, args: &SucceedModuleArgs<'_>) -> rspack_error::Result<()> {
if self.is_hook_disabled(&Hook::SucceedModule) {
return Ok(());
}
let js_module = args
.module
.to_js_module()
.expect("Failed to convert module to JsModule");
self
Expand Down
78 changes: 75 additions & 3 deletions crates/rspack_core/src/compiler/compilation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,9 @@ use rspack_identifier::{Identifiable, IdentifierMap, IdentifierSet};
use rspack_sources::{BoxSource, CachedSource, SourceExt};
use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet, FxHasher};
use swc_core::ecma::ast::ModuleItem;
use tokio::sync::mpsc::error::TryRecvError;
use tokio::sync::mpsc::{
error::TryRecvError, unbounded_channel, UnboundedReceiver, UnboundedSender,
};
use tracing::instrument;

use super::{
Expand All @@ -45,8 +47,8 @@ use crate::{
FactorizeTask, FactorizeTaskResult, Filename, Logger, Module, ModuleGraph, ModuleIdentifier,
ModuleProfile, ModuleType, PathData, ProcessAssetsArgs, ProcessDependenciesQueue,
ProcessDependenciesResult, ProcessDependenciesTask, RenderManifestArgs, Resolve, ResolverFactory,
RuntimeGlobals, RuntimeModule, RuntimeSpec, SharedPluginDriver, SourceType, Stats, TaskResult,
WorkerTask,
RuntimeGlobals, RuntimeModule, RuntimeSpec, SharedPluginDriver, SourceType, Stats,
SucceedModuleArgs, TaskResult, WorkerTask,
};
use crate::{tree_shaking::visitor::OptimizeAnalyzeResult, Context};

Expand Down Expand Up @@ -110,6 +112,11 @@ pub struct Compilation {
pub build_dependencies: IndexSet<PathBuf, BuildHasherDefault<FxHasher>>,
pub side_effects_free_modules: IdentifierSet,
pub module_item_map: IdentifierMap<Vec<ModuleItem>>,

pub factorize_queue: QueueHandler<FactorizeTask>,
pub add_queue: QueueHandler<AddTask>,
pub build_queue: QueueHandler<BuildTask>,
pub process_dependencies_queue: QueueHandler<ProcessDependenciesTask>,
}

impl Compilation {
Expand Down Expand Up @@ -170,6 +177,11 @@ impl Compilation {
side_effects_free_modules: IdentifierSet::default(),
module_item_map: IdentifierMap::default(),
include_module_ids: IdentifierSet::default(),

factorize_queue: QueueHandler::new(),
add_queue: QueueHandler::new(),
build_queue: QueueHandler::new(),
process_dependencies_queue: QueueHandler::new(),
}
}

Expand Down Expand Up @@ -431,10 +443,12 @@ impl Compilation {
let mut active_task_count = 0usize;
let is_expected_shutdown = Arc::new(AtomicBool::new(false));
let (result_tx, mut result_rx) = tokio::sync::mpsc::unbounded_channel::<Result<TaskResult>>();

let mut factorize_queue = FactorizeQueue::new();
let mut add_queue = AddQueue::new();
let mut build_queue = BuildQueue::new();
let mut process_dependencies_queue = ProcessDependenciesQueue::new();

let mut make_failed_dependencies: HashSet<BuildDependency> = HashSet::default();
let mut make_failed_module = HashSet::default();
let mut errored = None;
Expand Down Expand Up @@ -489,6 +503,19 @@ impl Compilation {
}

tokio::task::block_in_place(|| loop {
if let Ok(task) = self.factorize_queue.try_recv() {
factorize_queue.add_task(task);
}
if let Ok(task) = self.add_queue.try_recv() {
add_queue.add_task(task);
}
if let Ok(task) = self.build_queue.try_recv() {
build_queue.add_task(task);
}
if let Ok(task) = self.process_dependencies_queue.try_recv() {
process_dependencies_queue.add_task(task);
}

let start = factorize_time.start();
while let Some(task) = factorize_queue.get_task() {
tokio::spawn({
Expand Down Expand Up @@ -724,6 +751,18 @@ impl Compilation {
from_cache,
} = task_result;

tokio::runtime::Handle::current().block_on(async {
self
.plugin_driver
.succeed_module(&SucceedModuleArgs {
module: &module,
build_result: &build_result,
compilation: self,
})
.await
.unwrap_or_else(|e| panic!("Run succeed_module hook failed: {}", e));
});

if let Some(counter) = &mut build_cache_counter {
if from_cache {
counter.hit();
Expand Down Expand Up @@ -1844,3 +1883,36 @@ impl AssetInfo {
pub struct AssetInfoRelated {
pub source_map: Option<String>,
}

#[derive(Debug)]
pub struct QueueHandler<T>(UnboundedSender<T>, UnboundedReceiver<T>);

impl<T> Default for QueueHandler<T> {
fn default() -> Self {
let (tx, rx) = unbounded_channel();
Self(tx, rx)
}
}

impl<T> QueueHandler<T> {
pub fn new() -> Self {
Self::default()
}

pub fn add(&self, task: T) -> Result<()> {
self.0.send(task).map_err(|reason| {
rspack_error::Error::InternalError(rspack_error::InternalError::new(
format!("Failed to add queue task, reason: {}", reason),
Severity::Error,
))
})
}

pub fn get_tx(&self) -> UnboundedSender<T> {
self.0.clone()
}

pub fn try_recv(&mut self) -> std::result::Result<T, TryRecvError> {
self.1.try_recv()
}
}
9 changes: 4 additions & 5 deletions crates/rspack_core/src/compiler/queue.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ pub trait WorkerTask {
async fn run(self) -> Result<TaskResult>;
}

#[derive(Debug)]
pub struct FactorizeTask {
pub original_module_identifier: Option<ModuleIdentifier>,
pub original_module_context: Option<Box<Context>>,
Expand Down Expand Up @@ -170,6 +171,7 @@ impl WorkerTask for FactorizeTask {

pub type FactorizeQueue = WorkerQueue<FactorizeTask>;

#[derive(Debug)]
pub struct AddTask {
pub original_module_identifier: Option<ModuleIdentifier>,
pub module: Box<dyn Module>,
Expand Down Expand Up @@ -256,6 +258,7 @@ fn set_resolved_module(

pub type AddQueue = WorkerQueue<AddTask>;

#[derive(Debug)]
pub struct BuildTask {
pub module: Box<dyn Module>,
pub resolver_factory: Arc<ResolverFactory>,
Expand Down Expand Up @@ -308,11 +311,6 @@ impl WorkerTask for BuildTask {
})
.await;

plugin_driver
.succeed_module(&**module)
.await
.unwrap_or_else(|e| panic!("Run succeed_module hook failed: {}", e));

result.map(|t| (t, module))
})
.await
Expand Down Expand Up @@ -345,6 +343,7 @@ impl WorkerTask for BuildTask {

pub type BuildQueue = WorkerQueue<BuildTask>;

#[derive(Debug)]
pub struct ProcessDependenciesTask {
pub original_module_identifier: ModuleIdentifier,
pub dependencies: Vec<DependencyId>,
Expand Down
3 changes: 3 additions & 0 deletions crates/rspack_core/src/dependency/dependency_type.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,8 @@ pub enum DependencyType {
RequireResolve,
/// wasm import
WasmImport,
/// import module
LoaderImport,
/// wasm export import
WasmExportImported,
/// static exports
Expand Down Expand Up @@ -99,6 +101,7 @@ impl DependencyType {
DependencyType::RequireContext => Cow::Borrowed("require.context"),
DependencyType::RequireResolve => Cow::Borrowed("require.resolve"),
DependencyType::WasmImport => Cow::Borrowed("wasm import"),
DependencyType::LoaderImport => Cow::Borrowed("loader import"),
DependencyType::WasmExportImported => Cow::Borrowed("wasm export imported"),
DependencyType::StaticExports => Cow::Borrowed("static exports"),
DependencyType::Custom(ty) => Cow::Owned(format!("custom {ty}")),
Expand Down
49 changes: 49 additions & 0 deletions crates/rspack_core/src/dependency/loader_import_dependency.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
use crate::{
AsContextDependency, AsDependencyTemplate, Dependency, DependencyId, ModuleDependency,
};

#[derive(Debug, Hash, PartialEq, Eq, Clone)]
pub struct LoaderImportDependency {
request: String,
id: DependencyId,
}

impl LoaderImportDependency {
pub fn new(request: String) -> Self {
Self {
request,
id: DependencyId::new(),
}
}
}

impl AsDependencyTemplate for LoaderImportDependency {}
impl AsContextDependency for LoaderImportDependency {}

impl Dependency for LoaderImportDependency {
fn dependency_debug_name(&self) -> &'static str {
"LoaderImportDependency"
}

fn id(&self) -> &crate::DependencyId {
&self.id
}

fn dependency_type(&self) -> &crate::DependencyType {
&crate::DependencyType::LoaderImport
}
}

impl ModuleDependency for LoaderImportDependency {
fn request(&self) -> &str {
&self.request
}

fn user_request(&self) -> &str {
&self.request
}

fn set_request(&mut self, request: String) {
self.request = request
}
}
1 change: 1 addition & 0 deletions crates/rspack_core/src/normal_module_factory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -636,6 +636,7 @@ impl NormalModuleFactory {
.plugin_driver
.create_module(NormalModuleCreateData {
dependency_type: data.dependency.dependency_type().clone(),
dependency_id: *data.dependency.id(),
})
.await?
{
Expand Down
5 changes: 3 additions & 2 deletions crates/rspack_core/src/plugin/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,8 @@ use crate::{
NormalModule, NormalModuleAfterResolveArgs, NormalModuleBeforeResolveArgs,
NormalModuleCreateData, NormalModuleFactoryContext, OptimizeChunksArgs, ParserAndGenerator,
PluginContext, ProcessAssetsArgs, RenderArgs, RenderChunkArgs, RenderManifestArgs,
RenderModuleContentArgs, RenderStartupArgs, Resolver, SourceType, ThisCompilationArgs,
RenderModuleContentArgs, RenderStartupArgs, Resolver, SourceType, SucceedModuleArgs,
ThisCompilationArgs,
};

// use anyhow::{Context, Result};
Expand Down Expand Up @@ -435,7 +436,7 @@ pub trait Plugin: Debug + Send + Sync {
Ok(())
}

async fn succeed_module(&self, _module: &dyn Module) -> Result<()> {
async fn succeed_module(&self, _args: &SucceedModuleArgs<'_>) -> Result<()> {
Ok(())
}

Expand Down
13 changes: 10 additions & 3 deletions crates/rspack_core/src/plugin/args.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@ use rspack_sources::BoxSource;
use rustc_hash::FxHashSet as HashSet;

use crate::{
Chunk, ChunkInitFragments, ChunkUkey, Compilation, Context, DependencyCategory, DependencyType,
ErrorSpan, FactoryMeta, ModuleDependency, ModuleGraphModule, ModuleIdentifier, Resolve,
RuntimeGlobals, SharedPluginDriver, Stats,
BoxModule, BuildResult, Chunk, ChunkInitFragments, ChunkUkey, Compilation, Context,
DependencyCategory, DependencyId, DependencyType, ErrorSpan, FactoryMeta, ModuleDependency,
ModuleGraphModule, ModuleIdentifier, Resolve, RuntimeGlobals, SharedPluginDriver, Stats,
};

#[derive(Debug)]
Expand Down Expand Up @@ -84,6 +84,7 @@ pub struct FactorizeArgs<'me> {
#[derive(Debug, Clone)]
pub struct NormalModuleCreateData {
pub dependency_type: DependencyType,
pub dependency_id: DependencyId,
}

#[derive(Debug, Clone)]
Expand Down Expand Up @@ -247,3 +248,9 @@ impl<'me> JsChunkHashArgs<'me> {
.expect("chunk should exist in chunk_by_ukey")
}
}

pub struct SucceedModuleArgs<'me> {
pub module: &'me BoxModule,
pub build_result: &'me BuildResult,
pub compilation: &'me Compilation,
}
6 changes: 3 additions & 3 deletions crates/rspack_core/src/plugin/plugin_driver.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ use crate::{
PluginRenderManifestHookOutput, PluginRenderModuleContentOutput, PluginRenderStartupHookOutput,
PluginThisCompilationHookOutput, ProcessAssetsArgs, RenderArgs, RenderChunkArgs,
RenderManifestArgs, RenderModuleContentArgs, RenderStartupArgs, Resolver, ResolverFactory, Stats,
ThisCompilationArgs,
SucceedModuleArgs, ThisCompilationArgs,
};

pub struct PluginDriver {
Expand Down Expand Up @@ -589,9 +589,9 @@ impl PluginDriver {
}

#[instrument(name = "plugin:succeed_module", skip_all)]
pub async fn succeed_module(&self, module: &dyn Module) -> Result<()> {
pub async fn succeed_module(&self, args: &SucceedModuleArgs<'_>) -> Result<()> {
for plugin in &self.plugins {
plugin.succeed_module(module).await?;
plugin.succeed_module(args).await?;
}
Ok(())
}
Expand Down
4 changes: 4 additions & 0 deletions crates/rspack_error/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,10 @@ impl<T: std::fmt::Debug> TWithDiagnosticArray<T> {
pub fn split_into_parts(self) -> (T, Vec<Diagnostic>) {
(self.inner, self.diagnostic)
}

pub fn get(&self) -> &T {
&self.inner
}
}

impl<T: Clone + std::fmt::Debug> Clone for TWithDiagnosticArray<T> {
Expand Down
2 changes: 2 additions & 0 deletions crates/rspack_util/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ version = "0.1.0"

[dependencies]
concat-string = { workspace = true }
dashmap = { workspace = true }
once_cell = { workspace = true }
regex = { workspace = true }
rustc-hash = { workspace = true }
sugar_path = { workspace = true }
7 changes: 7 additions & 0 deletions crates/rspack_util/src/fx_dashmap.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
use std::hash::BuildHasherDefault;

use dashmap::{DashMap, DashSet};
use rustc_hash::FxHasher;

pub type FxDashMap<K, V> = DashMap<K, V, BuildHasherDefault<FxHasher>>;
pub type FxDashSet<V> = DashSet<V, BuildHasherDefault<FxHasher>>;
1 change: 1 addition & 0 deletions crates/rspack_util/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
use std::future::Future;
pub mod comparators;
pub mod ext;
pub mod fx_dashmap;
pub mod identifier;
pub mod number_hash;
pub mod swc;
Expand Down

0 comments on commit 991cb90

Please sign in to comment.