feat(rust/dev): add DevEngine#close (#6376)

hyf0 · hyf0 · commit 0f920ceb39c5 · 2025-09-30T05:36:13.000Z
diff --git a/crates/rolldown/examples/dev.rs b/crates/rolldown/examples/dev.rs
@@ -28,5 +28,5 @@ async fn main() {
   )
   .unwrap();
   dev_engine.run().await.unwrap();
-  dev_engine.wait_for_close().await;
+  dev_engine.wait_for_build_driver_service_close().await.unwrap();
 }
diff --git a/crates/rolldown/src/dev/build_driver_service.rs b/crates/rolldown/src/dev/build_driver_service.rs
@@ -10,6 +10,7 @@ use crate::dev::{
 pub enum BuildMessage {
   WatchEvent(FileChangeResult),
   BuildFinish,
+  Close,
 }
 
 pub type BuildChannelTx = UnboundedSender<BuildMessage>;
@@ -74,6 +75,9 @@ impl BuildDriverService {
             .await
             .expect("FIXME: should handle this error");
         }
+        BuildMessage::Close => {
+          break;
+        }
       }
     }
   }
diff --git a/crates/rolldown/src/dev/dev_engine.rs b/crates/rolldown/src/dev/dev_engine.rs
@@ -1,4 +1,9 @@
-use std::{collections::VecDeque, ops::Deref, path::PathBuf, sync::Arc};
+use std::{
+  collections::VecDeque,
+  ops::Deref,
+  path::PathBuf,
+  sync::{Arc, atomic::AtomicBool},
+};
 
 use arcstr::ArcStr;
 use futures::{FutureExt, future::Shared};
@@ -32,9 +37,10 @@ pub struct DevEngine {
   build_driver: SharedBuildDriver,
   watcher: Mutex<DynWatcher>,
   watched_files: FxDashSet<ArcStr>,
-  watch_service_state: Mutex<BuildDriverServiceState>,
+  build_driver_service_state: Mutex<BuildDriverServiceState>,
   ctx: SharedDevContext,
   pub clients: SharedClients,
+  is_closed: AtomicBool,
 }
 
 impl DevEngine {
@@ -128,17 +134,18 @@ impl DevEngine {
       build_driver,
       watcher: Mutex::new(watcher),
       watched_files: FxDashSet::default(),
-      watch_service_state: Mutex::new(BuildDriverServiceState {
+      build_driver_service_state: Mutex::new(BuildDriverServiceState {
         service: Some(build_driver_service),
         handle: None,
       }),
       ctx,
       clients,
+      is_closed: AtomicBool::new(false),
     })
   }
 
   pub async fn run(&self) -> BuildResult<()> {
-    let mut build_service_state = self.watch_service_state.lock().await;
+    let mut build_service_state = self.build_driver_service_state.lock().await;
 
     if build_service_state.service.is_none() {
       // The watcher service is already running.
@@ -174,25 +181,58 @@ impl DevEngine {
     Ok(())
   }
 
-  pub async fn wait_for_close(&self) {
-    let watch_service_state = self.watch_service_state.lock().await;
-    if let Some(watcher_service_handle) = watch_service_state.handle.clone() {
-      watcher_service_handle.await;
+  pub async fn wait_for_build_driver_service_close(&self) -> BuildResult<()> {
+    self.create_error_if_closed()?;
+    let service_state = self.build_driver_service_state.lock().await;
+    if let Some(service_handle) = service_state.handle.clone() {
+      service_handle.await;
     }
+    Ok(())
   }
 
-  pub async fn ensure_current_build_finish(&self) {
+  pub async fn ensure_current_build_finish(&self) -> BuildResult<()> {
+    self.create_error_if_closed()?;
     self.ctx.ensure_current_build_finish().await;
+    Ok(())
   }
 
   pub async fn invalidate(
     &self,
     caller: String,
     first_invalidated_by: Option<String>,
   ) -> BuildResult<Vec<ClientHmrUpdate>> {
+    self.create_error_if_closed()?;
     self.build_driver.invalidate(caller, first_invalidated_by).await
   }
 
+  pub async fn close(&self) -> BuildResult<()> {
+    if self.is_closed.swap(true, std::sync::atomic::Ordering::SeqCst) {
+      return Ok(());
+    }
+
+    // Send close message to build driver service
+    if let Err(_e) = self.ctx.build_channel_tx.send(BuildMessage::Close) {
+      // If `BuildMessage::Close` is not sent, it means the channel is closed due to error or what.
+      // It's ok to ignore the send error.
+    }
+
+    // Clean up watcher
+    let watcher =
+      std::mem::replace(&mut *self.watcher.lock().await, NoopWatcher.into_dyn_watcher());
+    std::mem::drop(watcher);
+
+    // Close the bundler
+    let mut bundler = self.build_driver.bundler.lock().await;
+    bundler.close().await?;
+
+    // Wait for build driver service to close
+    let service_state = self.build_driver_service_state.lock().await;
+    if let Some(service_handle) = service_state.handle.clone() {
+      service_handle.await;
+    }
+    Ok(())
+  }
+
   /// For testing purpose.
   pub async fn ensure_task_with_changed_files(&self, changed_files: FxIndexSet<PathBuf>) {
     self.build_driver.handle_file_changes(changed_files).await;
@@ -211,6 +251,17 @@ impl DevEngine {
     }
     self.clients.insert("test".to_string(), client_session);
   }
+
+  pub fn is_closed(&self) -> bool {
+    self.is_closed.load(std::sync::atomic::Ordering::SeqCst)
+  }
+
+  fn create_error_if_closed(&self) -> BuildResult<()> {
+    if self.is_closed.load(std::sync::atomic::Ordering::SeqCst) {
+      Err(anyhow::anyhow!("Dev engine is closed"))?;
+    }
+    Ok(())
+  }
 }
 
 impl Deref for DevEngine {
diff --git a/crates/rolldown_binding/src/binding_dev_engine.rs b/crates/rolldown_binding/src/binding_dev_engine.rs
@@ -96,7 +96,11 @@ impl BindingDevEngine {
 
   #[napi]
   pub async fn ensure_current_build_finish(&self) -> napi::Result<()> {
-    self.inner.ensure_current_build_finish().await;
+    self
+      .inner
+      .ensure_current_build_finish()
+      .await
+      .map_err(|_e| napi::Error::from_reason("Failed to ensure current build finish"))?;
     Ok(())
   }
 

Original file line number	Diff line number	Diff line change
`@@ -28,5 +28,5 @@ async fn main() {`
`28`	`28`	`)`
`29`	`29`	`.unwrap();`
`30`	`30`	`dev_engine.run().await.unwrap();`
`31`		`- dev_engine.wait_for_close().await;`
	`31`	`+ dev_engine.wait_for_build_driver_service_close().await.unwrap();`
`32`	`32`	`}`
Original file line number	Diff line number	Diff line change
`@@ -10,6 +10,7 @@ use crate::dev::{`
`10`	`10`	`pub enum BuildMessage {`
`11`	`11`	`WatchEvent(FileChangeResult),`
`12`	`12`	`BuildFinish,`
	`13`	`+ Close,`
`13`	`14`	`}`
`14`	`15`
`15`	`16`	`pub type BuildChannelTx = UnboundedSender<BuildMessage>;`
`@@ -74,6 +75,9 @@ impl BuildDriverService {`
`74`	`75`	`.await`
`75`	`76`	`.expect("FIXME: should handle this error");`
`76`	`77`	`}`
	`78`	`+ BuildMessage::Close => {`
	`79`	`+ break;`
	`80`	`+ }`
`77`	`81`	`}`
`78`	`82`	`}`
`79`	`83`	`}`