Skip to content

Commit

Permalink
👀 hot shader reloading (#655)
Browse files Browse the repository at this point in the history
* Update builder to use a shared validation method

* Add the error for using print_metadata and watching

We cannot use print_metadata with watching
because print_metadata only makes sense
in build scripts, but watching does not
and would instead stall the script

* Add the initial implementation of watching

* Make hot reloading work in the wgpu example

* Attempt to address CI failures

* Add exception for notify CC0-1.0 license

* Address review comments

Co-authored-by: khyperia <github@khyperia.com>
  • Loading branch information
DJMcNab and khyperia authored Jun 9, 2021
1 parent 0410fc5 commit 3bbe963
Show file tree
Hide file tree
Showing 9 changed files with 370 additions and 89 deletions.
73 changes: 71 additions & 2 deletions Cargo.lock

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

3 changes: 3 additions & 0 deletions crates/spirv-builder/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ license = "MIT OR Apache-2.0"
default = ["use-compiled-tools"]
use-installed-tools = ["rustc_codegen_spirv/use-installed-tools"]
use-compiled-tools = ["rustc_codegen_spirv/use-compiled-tools"]
watch = ["notify"]

[dependencies]
memchr = "2.3"
Expand All @@ -18,3 +19,5 @@ serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
# See comment in lib.rs invoke_rustc for why this is here
rustc_codegen_spirv = { path = "../rustc_codegen_spirv", default-features = false }

notify = { version = "5.0.0-pre.10", optional = true }
76 changes: 54 additions & 22 deletions crates/spirv-builder/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,8 @@
#![allow()]

mod depfile;
#[cfg(feature = "watch")]
mod watch;

use raw_string::{RawStr, RawString};
use serde::Deserialize;
Expand All @@ -71,10 +73,12 @@ pub use rustc_codegen_spirv::rspirv::spirv::Capability;
pub use rustc_codegen_spirv::{CompileResult, ModuleResult};

#[derive(Debug)]
#[non_exhaustive]
pub enum SpirvBuilderError {
CratePathDoesntExist(PathBuf),
BuildFailed,
MultiModuleWithPrintMetadata,
WatchWithPrintMetadata,
MetadataFileMissing(std::io::Error),
MetadataFileMalformed(serde_json::Error),
}
Expand All @@ -89,6 +93,9 @@ impl fmt::Display for SpirvBuilderError {
SpirvBuilderError::MultiModuleWithPrintMetadata => f.write_str(
"Multi-module build cannot be used with print_metadata = MetadataPrintout::Full",
),
SpirvBuilderError::WatchWithPrintMetadata => {
f.write_str("Watching within build scripts will prevent build completion")
}
SpirvBuilderError::MetadataFileMissing(_) => {
f.write_str("Multi-module metadata file missing")
}
Expand Down Expand Up @@ -244,22 +251,47 @@ impl SpirvBuilder {

/// Builds the module. If `print_metadata` is [`MetadataPrintout::Full`], you usually don't have to inspect the path
/// in the result, as the environment variable for the path to the module will already be set.
pub fn build(self) -> Result<CompileResult, SpirvBuilderError> {
pub fn build(mut self) -> Result<CompileResult, SpirvBuilderError> {
self.validate_running_conditions()?;
let metadata_file = invoke_rustc(&self)?;
match self.print_metadata {
MetadataPrintout::Full | MetadataPrintout::DependencyOnly => {
leaf_deps(&metadata_file, |artifact| {
println!("cargo:rerun-if-changed={}", artifact)
})
// Close enough
.map_err(SpirvBuilderError::MetadataFileMissing)?;
}
MetadataPrintout::None => (),
}
let metadata = self.parse_metadata_file(&metadata_file)?;

Ok(metadata)
}

pub(crate) fn validate_running_conditions(&mut self) -> Result<(), SpirvBuilderError> {
if (self.print_metadata == MetadataPrintout::Full) && self.multimodule {
return Err(SpirvBuilderError::MultiModuleWithPrintMetadata);
}
if !self.path_to_crate.is_dir() {
return Err(SpirvBuilderError::CratePathDoesntExist(self.path_to_crate));
return Err(SpirvBuilderError::CratePathDoesntExist(std::mem::take(
&mut self.path_to_crate,
)));
}
let metadata_file = invoke_rustc(&self)?;
let metadata_contents =
File::open(&metadata_file).map_err(SpirvBuilderError::MetadataFileMissing)?;
Ok(())
}

pub(crate) fn parse_metadata_file(
&self,
at: &Path,
) -> Result<CompileResult, SpirvBuilderError> {
let metadata_contents = File::open(&at).map_err(SpirvBuilderError::MetadataFileMissing)?;
let metadata: CompileResult = serde_json::from_reader(BufReader::new(metadata_contents))
.map_err(SpirvBuilderError::MetadataFileMalformed)?;
match &metadata.module {
ModuleResult::SingleModule(spirv_module) => {
assert!(!self.multimodule);
let env_var = metadata_file.file_name().unwrap().to_str().unwrap();
let env_var = at.file_name().unwrap().to_str().unwrap();
if self.print_metadata == MetadataPrintout::Full {
println!("cargo:rustc-env={}={}", env_var, spirv_module.display());
}
Expand Down Expand Up @@ -424,13 +456,8 @@ fn invoke_rustc(builder: &SpirvBuilder) -> Result<PathBuf, SpirvBuilderError> {
// that ended up on stdout instead of stderr.
let stdout = String::from_utf8(build.stdout).unwrap();
let artifact = get_last_artifact(&stdout);

if build.status.success() {
match builder.print_metadata {
MetadataPrintout::Full | MetadataPrintout::DependencyOnly => print_deps_of(&artifact),
MetadataPrintout::None => (),
}
Ok(artifact)
Ok(artifact.expect("Artifact created when compilation succeeded"))
} else {
Err(SpirvBuilderError::BuildFailed)
}
Expand All @@ -442,7 +469,7 @@ struct RustcOutput {
filenames: Option<Vec<String>>,
}

fn get_last_artifact(out: &str) -> PathBuf {
fn get_last_artifact(out: &str) -> Option<PathBuf> {
let last = out
.lines()
.filter_map(|line| match serde_json::from_str::<RustcOutput>(line) {
Expand All @@ -462,28 +489,33 @@ fn get_last_artifact(out: &str) -> PathBuf {
.unwrap()
.into_iter()
.filter(|v| v.ends_with(".spv"));
let filename = filenames.next().expect("Crate had no .spv artifacts");
let filename = filenames.next()?;
assert_eq!(filenames.next(), None, "Crate had multiple .spv artifacts");
filename.into()
Some(filename.into())
}

fn print_deps_of(artifact: &Path) {
/// Internally iterate through the leaf dependencies of the artifact at `artifact`
fn leaf_deps(artifact: &Path, mut handle: impl FnMut(&RawStr)) -> std::io::Result<()> {
let deps_file = artifact.with_extension("d");
let mut deps_map = HashMap::new();
depfile::read_deps_file(&deps_file, |item, deps| {
deps_map.insert(item, deps);
Ok(())
})
.expect("Could not read dep file");
fn recurse(map: &HashMap<RawString, Vec<RawString>>, artifact: &RawStr) {
})?;
fn recurse(
map: &HashMap<RawString, Vec<RawString>>,
artifact: &RawStr,
handle: &mut impl FnMut(&RawStr),
) {
match map.get(artifact) {
Some(entries) => {
for entry in entries {
recurse(map, entry)
recurse(map, entry, handle)
}
}
None => println!("cargo:rerun-if-changed={}", artifact),
None => handle(artifact),
}
}
recurse(&deps_map, artifact.to_str().unwrap().into());
recurse(&deps_map, artifact.to_str().unwrap().into(), &mut handle);
Ok(())
}
104 changes: 104 additions & 0 deletions crates/spirv-builder/src/watch.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
use std::{collections::HashSet, sync::mpsc::sync_channel};

use notify::{Event, RecursiveMode, Watcher};
use rustc_codegen_spirv::CompileResult;

use crate::{leaf_deps, SpirvBuilder, SpirvBuilderError};

impl SpirvBuilder {
/// Watches the module for changes using [`notify`](https://crates.io/crates/notify).
///
/// This is a blocking operation, wand should never return in the happy path
pub fn watch(
mut self,
on_compilation_finishes: impl Fn(CompileResult),
) -> Result<(), SpirvBuilderError> {
self.validate_running_conditions()?;
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 = match metadata_result {
Ok(path) => path,
Err(_) => {
let (tx, rx) = sync_channel(0);
// Fall back to watching from the crate root if the inital compilation fails
let mut watcher =
notify::immediate_watcher(move |event: notify::Result<Event>| match event {
Ok(e) => match e.kind {
notify::EventKind::Access(_) => (),
notify::EventKind::Any
| notify::EventKind::Create(_)
| notify::EventKind::Modify(_)
| notify::EventKind::Remove(_)
| notify::EventKind::Other => {
let _ = tx.try_send(());
}
},
Err(e) => println!("notify error: {:?}", e),
})
.expect("Could create watcher");
// This is likely to notice changes in the `target` dir, however, given that `cargo watch` doesn't seem to handle that,
watcher
.watch(&self.path_to_crate, RecursiveMode::Recursive)
.expect("Could watch crate root");
loop {
rx.recv().expect("Watcher still alive");
let metadata_file = crate::invoke_rustc(&self);
if let Ok(f) = metadata_file {
break f;
}
}
}
};
let metadata = self.parse_metadata_file(&metadata_file)?;
on_compilation_finishes(metadata);
let mut watched_paths = HashSet::new();
let (tx, rx) = sync_channel(0);
let mut watcher =
notify::immediate_watcher(move |event: notify::Result<Event>| match event {
Ok(e) => match e.kind {
notify::EventKind::Access(_) => (),
notify::EventKind::Any
| notify::EventKind::Create(_)
| notify::EventKind::Modify(_)
| notify::EventKind::Remove(_)
| notify::EventKind::Other => {
let _ = tx.try_send(());
}
},
Err(e) => println!("notify error: {:?}", e),
})
.expect("Could create watcher");
leaf_deps(&metadata_file, |it| {
let path = it.to_path().unwrap();
if watched_paths.insert(path.to_owned()) {
watcher
.watch(it.to_path().unwrap(), RecursiveMode::NonRecursive)
.expect("Cargo dependencies are valid files");
}
})
.expect("Could read dependencies file");
loop {
rx.recv().expect("Watcher still alive");
let metadata_result = crate::invoke_rustc(&self);
if let Ok(file) = metadata_result {
// We can bubble this error up because it's an internal error (e.g. rustc_codegen_spirv's version of CompileResult is somehow out of sync)
let metadata = self.parse_metadata_file(&file)?;

leaf_deps(&file, |it| {
let path = it.to_path().unwrap();
if watched_paths.insert(path.to_owned()) {
watcher
.watch(it.to_path().unwrap(), RecursiveMode::NonRecursive)
.expect("Cargo dependencies are valid files");
}
})
.expect("Could read dependencies file");

on_compilation_finishes(metadata);
}
}
}
}
Loading

0 comments on commit 3bbe963

Please sign in to comment.