Redesign spirv-builder's watch API

Author: tuguzT

crates/spirv-builder/src/lib.rs

@@ -92,7 +92,7 @@ pub use rustc_codegen_spirv_types::Capability;
 pub use rustc_codegen_spirv_types::{CompileResult, ModuleResult};
 
 #[cfg(feature = "watch")]
-pub use self::watch::Watch;
+pub use self::watch::{SpirvWatcher, SpirvWatcherError};
 
 #[cfg(feature = "include-target-specs")]
 pub use rustc_codegen_spirv_target_specs::TARGET_SPEC_DIR_PATH;
@@ -125,14 +125,15 @@ pub enum SpirvBuilderError {
     BuildFailed,
     #[error("multi-module build cannot be used with print_metadata = MetadataPrintout::Full")]
     MultiModuleWithPrintMetadata,
-    #[error("watching within build scripts will prevent build completion")]
-    WatchWithPrintMetadata,
     #[error("multi-module metadata file missing")]
     MetadataFileMissing(#[from] std::io::Error),
     #[error("unable to parse multi-module metadata file")]
     MetadataFileMalformed(#[from] serde_json::Error),
     #[error("cargo metadata error")]
     CargoMetadata(#[from] cargo_metadata::Error),
+    #[cfg(feature = "watch")]
+    #[error(transparent)]
+    WatchFailed(#[from] SpirvWatcherError),
 }
 
 const SPIRV_TARGET_PREFIX: &str = "spirv-unknown-";
@@ -518,6 +519,12 @@ impl Default for SpirvBuilder {
     }
 }
 
+impl AsRef<SpirvBuilder> for SpirvBuilder {
+    fn as_ref(&self) -> &SpirvBuilder {
+        self
+    }
+}
+
 impl SpirvBuilder {
     pub fn new(path_to_crate: impl AsRef<Path>, target: impl Into<String>) -> Self {
         Self {

crates/spirv-builder/src/watch.rs

@@ -1,153 +1,149 @@
-use std::convert::Infallible;
 use std::path::{Path, PathBuf};
 use std::sync::mpsc::Receiver;
-use std::thread::JoinHandle;
 use std::{collections::HashSet, sync::mpsc::sync_channel};
 
-use notify::{Event, RecommendedWatcher, RecursiveMode, Watcher as _};
+use notify::{Event, RecommendedWatcher, RecursiveMode, Watcher};
 use rustc_codegen_spirv_types::CompileResult;
 
 use crate::{SpirvBuilder, SpirvBuilderError, leaf_deps};
 
 impl SpirvBuilder {
-    /// Watches the module for changes using [`notify`], rebuilding it upon changes.
-    ///
-    /// Calls `on_compilation_finishes` after each successful compilation.
-    /// The second `Option<AcceptFirstCompile<T>>` param allows you to return some `T`
-    /// on the first compile, which is then returned by this function
-    /// in pair with [`JoinHandle`] to the watching thread.
-    pub fn watch<T>(
-        &self,
-        mut on_compilation_finishes: impl FnMut(CompileResult, Option<AcceptFirstCompile<'_, T>>)
-        + Send
-        + 'static,
-    ) -> Result<Watch<T>, SpirvBuilderError> {
-        let path_to_crate = self
-            .path_to_crate
-            .as_ref()
-            .ok_or(SpirvBuilderError::MissingCratePath)?;
-        if !matches!(self.print_metadata, crate::MetadataPrintout::None) {
-            return Err(SpirvBuilderError::WatchWithPrintMetadata);
-        }
-
-        let metadata_result = crate::invoke_rustc(self);
-        // Load the dependencies of the thing
-        let metadata_file = if let Ok(path) = metadata_result {
-            path
-        } else {
-            // Fall back to watching from the crate root if the initial compilation fails
-            // This is likely to notice changes in the `target` dir, however, given that `cargo watch` doesn't seem to handle that,
-            let mut watcher = Watcher::new();
-            watcher
-                .watcher
-                .watch(path_to_crate, RecursiveMode::Recursive)
-                .expect("Could watch crate root");
-            loop {
-                watcher.recv();
-                let metadata_file = crate::invoke_rustc(self);
-                if let Ok(f) = metadata_file {
-                    break f;
-                }
-            }
-        };
-        let metadata = self.parse_metadata_file(&metadata_file)?;
-        let mut first_compile = None;
-        on_compilation_finishes(metadata, Some(AcceptFirstCompile(&mut first_compile)));
-
-        let builder = self.clone();
-        let watch_thread = std::thread::spawn(move || {
-            let mut watcher = Watcher::new();
-            watcher.watch_leaf_deps(&metadata_file);
-
-            loop {
-                watcher.recv();
-                let metadata_result = crate::invoke_rustc(&builder);
-                if let Ok(file) = metadata_result {
-                    let metadata = builder
-                        .parse_metadata_file(&file)
-                        .expect("Metadata file is correct");
-                    watcher.watch_leaf_deps(&metadata_file);
-                    on_compilation_finishes(metadata, None);
-                }
-            }
-        });
-
-        Ok(Watch {
-            first_compile,
-            watch_thread,
-        })
-    }
-}
-
-pub struct AcceptFirstCompile<'a, T>(&'a mut Option<T>);
-
-impl<'a, T> AcceptFirstCompile<'a, T> {
-    pub fn new(write: &'a mut Option<T>) -> Self {
-        Self(write)
-    }
-
-    pub fn submit(self, t: T) {
-        *self.0 = Some(t);
+    /// Watches the module for changes, rebuilding it upon them.
+    pub fn watch(&self) -> Result<SpirvWatcher<&Self>, SpirvBuilderError> {
+        SpirvWatcher::new(self)
     }
 }
 
-/// Result of [watching](SpirvBuilder::watch) a module for changes.
-#[must_use]
-#[non_exhaustive]
-pub struct Watch<T> {
-    /// Result of the first compile, if any.
-    pub first_compile: Option<T>,
-    /// Join handle to the watching thread.
-    ///
-    /// You can drop it to detach the watching thread,
-    /// or [`join()`](JoinHandle::join) it to block the current thread until shutdown of the program.
-    pub watch_thread: JoinHandle<Infallible>,
-}
-
-struct Watcher {
+#[derive(Debug)]
+pub struct SpirvWatcher<B> {
+    builder: B,
     watcher: RecommendedWatcher,
     rx: Receiver<()>,
+    watch_path: PathBuf,
     watched_paths: HashSet<PathBuf>,
+    first_result: bool,
 }
 
-impl Watcher {
-    fn new() -> Self {
+impl<B> SpirvWatcher<B>
+where
+    B: AsRef<SpirvBuilder>,
+{
+    fn new(as_builder: B) -> Result<Self, SpirvBuilderError> {
+        let builder = as_builder.as_ref();
+        let path_to_crate = builder
+            .path_to_crate
+            .as_ref()
+            .ok_or(SpirvBuilderError::MissingCratePath)?;
+        if !matches!(builder.print_metadata, crate::MetadataPrintout::None) {
+            return Err(SpirvWatcherError::WatchWithPrintMetadata.into());
+        }
+
         let (tx, rx) = sync_channel(0);
         let watcher =
-            notify::recommended_watcher(move |event: notify::Result<Event>| match event {
-                Ok(e) => match e.kind {
-                    notify::EventKind::Access(_) => (),
+            notify::recommended_watcher(move |result: notify::Result<Event>| match result {
+                Ok(event) => match event.kind {
                     notify::EventKind::Any
                     | notify::EventKind::Create(_)
                     | notify::EventKind::Modify(_)
                     | notify::EventKind::Remove(_)
                     | notify::EventKind::Other => {
-                        let _ = tx.try_send(());
+                        if let Err(err) = tx.try_send(()) {
+                            log::error!("send error: {err:?}");
+                        }
                     }
+                    notify::EventKind::Access(_) => {}
                 },
-                Err(e) => println!("notify error: {e:?}"),
+                Err(err) => log::error!("notify error: {err:?}"),
             })
-            .expect("Could create watcher");
-        Self {
+            .map_err(SpirvWatcherError::NotifyFailed)?;
+
+        Ok(Self {
+            watch_path: path_to_crate.clone(),
+            builder: as_builder,
             watcher,
             rx,
             watched_paths: HashSet::new(),
+            first_result: false,
+        })
+    }
+
+    pub fn recv(&mut self) -> Result<CompileResult, SpirvBuilderError> {
+        if !self.first_result {
+            return self.recv_first_result();
         }
+
+        self.rx.recv().expect("watcher should be alive");
+        let builder = self.builder.as_ref();
+        let metadata_file = crate::invoke_rustc(builder)?;
+        let result = builder.parse_metadata_file(&metadata_file)?;
+
+        self.watch_leaf_deps(&self.watch_path.clone())?;
+        Ok(result)
     }
 
-    fn watch_leaf_deps(&mut self, metadata_file: &Path) {
+    fn recv_first_result(&mut self) -> Result<CompileResult, SpirvBuilderError> {
+        let builder = self.builder.as_ref();
+        let metadata_file = match crate::invoke_rustc(builder) {
+            Ok(path) => path,
+            Err(err) => {
+                log::error!("{err}");
+
+                self.watcher
+                    .watch(&self.watch_path, RecursiveMode::Recursive)
+                    .map_err(SpirvWatcherError::NotifyFailed)?;
+                let path = loop {
+                    self.rx.recv().expect("watcher should be alive");
+                    match crate::invoke_rustc(builder) {
+                        Ok(path) => break path,
+                        Err(err) => log::error!("{err}"),
+                    }
+                };
+                self.watcher
+                    .unwatch(&self.watch_path)
+                    .map_err(SpirvWatcherError::NotifyFailed)?;
+                path
+            }
+        };
+        let result = builder.parse_metadata_file(&metadata_file)?;
+
+        self.watch_leaf_deps(&metadata_file)?;
+        self.watch_path = metadata_file;
+        self.first_result = true;
+        Ok(result)
+    }
+
+    fn watch_leaf_deps(&mut self, metadata_file: &Path) -> Result<(), SpirvBuilderError> {
         leaf_deps(metadata_file, |it| {
             let path = it.to_path().unwrap();
-            if self.watched_paths.insert(path.to_owned()) {
-                self.watcher
+            if self.watched_paths.insert(path.to_owned())
+                && let Err(err) = self
+                    .watcher
                     .watch(it.to_path().unwrap(), RecursiveMode::NonRecursive)
-                    .expect("Cargo dependencies are valid files");
+            {
+                log::error!("files of cargo dependencies are not valid: {err}");
             }
         })
-        .expect("Could read dependencies file");
+        .map_err(SpirvBuilderError::MetadataFileMissing)
     }
+}
 
-    fn recv(&self) {
-        self.rx.recv().expect("Watcher still alive");
+impl SpirvWatcher<&SpirvBuilder> {
+    pub fn forget_lifetime(self) -> SpirvWatcher<SpirvBuilder> {
+        SpirvWatcher {
+            builder: self.builder.clone(),
+            watcher: self.watcher,
+            rx: self.rx,
+            watch_path: self.watch_path,
+            watched_paths: self.watched_paths,
+            first_result: self.first_result,
+        }
     }
 }
+
+#[derive(Debug, thiserror::Error)]
+pub enum SpirvWatcherError {
+    #[error("watching within build scripts will prevent build completion")]
+    WatchWithPrintMetadata,
+    #[error("could not notify for changes: {0}")]
+    NotifyFailed(#[from] notify::Error),
+}

examples/runners/wgpu/src/lib.rs

@@ -177,18 +177,19 @@ fn maybe_watch(
         }
 
         if let Some(mut f) = on_watch {
-            builder
-                .watch(move |compile_result, accept| {
+            let mut watcher = builder.watch().unwrap();
+            let first_compile = watcher.recv().unwrap();
+
+            let mut thread_watcher = watcher.forget_lifetime();
+            std::thread::spawn(move || {
+                loop {
+                    let compile_result = thread_watcher.recv().unwrap();
                     let modules = handle_compile_result(compile_result);
-                    if let Some(accept) = accept {
-                        accept.submit(modules);
-                    } else {
-                        f(modules);
-                    }
-                })
-                .expect("Configuration is correct for watching")
-                .first_compile
-                .unwrap()
+                    f(modules);
+                }
+            });
+
+            handle_compile_result(first_compile)
         } else {
             handle_compile_result(builder.build().unwrap())
         }