Internal Asset Hot Reloading (#3966)

Adds "hot reloading" of internal assets, which is normally not possible because they are loaded using `include_str` / direct Asset collection access.

This is accomplished via the following:
* Add a new `debug_asset_server` feature flag
* When that feature flag is enabled, create a second App with a second AssetServer that points to a configured location (by default the `crates` folder). Plugins that want to add hot reloading support for their assets can call the new `app.add_debug_asset::<T>()` and `app.init_debug_asset_loader::<T>()` functions.
* Load "internal" assets using the new `load_internal_asset` macro. By default this is identical to the current "include_str + register in asset collection" approach. But if the `debug_asset_server` feature flag is enabled, it will also load the asset dynamically in the debug asset server using the file path. It will then set up a correlation between the "debug asset" and the "actual asset" by listening for asset change events.

This is an alternative to #3673. The goal was to keep the boilerplate and features flags to a minimum for bevy plugin authors, and allow them to home their shaders near relevant code. 

This is a draft because I haven't done _any_ quality control on this yet. I'll probably rename things and remove a bunch of unwraps. I just got it working and wanted to use it to start a conversation.

Fixes #3660

Author: cart

Cargo.toml

@@ -90,6 +90,9 @@ subpixel_glyph_atlas = ["bevy_internal/subpixel_glyph_atlas"]
 # Enable systems that allow for automated testing on CI
 bevy_ci_testing = ["bevy_internal/bevy_ci_testing"]
 
+# Enable the "debug asset server" for hot reloading internal assets
+debug_asset_server = ["bevy_internal/debug_asset_server"]
+
 [dependencies]
 bevy_dylib = { path = "crates/bevy_dylib", version = "0.6.0", default-features = false, optional = true }
 bevy_internal = { path = "crates/bevy_internal", version = "0.6.0", default-features = false }

crates/bevy_asset/Cargo.toml

@@ -11,6 +11,7 @@ keywords = ["bevy"]
 [features]
 default = []
 filesystem_watcher = ["notify"]
+debug_asset_server = ["filesystem_watcher"]
 
 [dependencies]
 # bevy

crates/bevy_asset/src/asset_server.rs

@@ -92,6 +92,10 @@ impl AssetServer {
         }
     }
 
+    pub fn asset_io(&self) -> &dyn AssetIo {
+        &*self.server.asset_io
+    }
+
     pub(crate) fn register_asset_type<T: Asset>(&self) -> Assets<T> {
         if self
             .server

crates/bevy_asset/src/assets.rs

@@ -259,9 +259,15 @@ impl<T: Asset> Assets<T> {
 /// [App] extension methods for adding new asset types
 pub trait AddAsset {
     fn add_asset<T>(&mut self) -> &mut Self
+    where
+        T: Asset;
+    fn add_debug_asset<T: Clone>(&mut self) -> &mut Self
     where
         T: Asset;
     fn init_asset_loader<T>(&mut self) -> &mut Self
+    where
+        T: AssetLoader + FromWorld;
+    fn init_debug_asset_loader<T>(&mut self) -> &mut Self
     where
         T: AssetLoader + FromWorld;
     fn add_asset_loader<T>(&mut self, loader: T) -> &mut Self
@@ -292,6 +298,23 @@ impl AddAsset for App {
             .add_event::<AssetEvent<T>>()
     }
 
+    fn add_debug_asset<T: Clone>(&mut self) -> &mut Self
+    where
+        T: Asset,
+    {
+        #[cfg(feature = "debug_asset_server")]
+        {
+            self.add_system(crate::debug_asset_server::sync_debug_assets::<T>);
+            let mut app = self
+                .world
+                .get_non_send_resource_mut::<crate::debug_asset_server::DebugAssetApp>()
+                .unwrap();
+            app.add_asset::<T>()
+                .init_resource::<crate::debug_asset_server::HandleMap<T>>();
+        }
+        self
+    }
+
     fn init_asset_loader<T>(&mut self) -> &mut Self
     where
         T: AssetLoader + FromWorld,
@@ -300,6 +323,21 @@ impl AddAsset for App {
         self.add_asset_loader(result)
     }
 
+    fn init_debug_asset_loader<T>(&mut self) -> &mut Self
+    where
+        T: AssetLoader + FromWorld,
+    {
+        #[cfg(feature = "debug_asset_server")]
+        {
+            let mut app = self
+                .world
+                .get_non_send_resource_mut::<crate::debug_asset_server::DebugAssetApp>()
+                .unwrap();
+            app.init_asset_loader::<T>();
+        }
+        self
+    }
+
     fn add_asset_loader<T>(&mut self, loader: T) -> &mut Self
     where
         T: AssetLoader,
@@ -312,6 +350,43 @@ impl AddAsset for App {
     }
 }
 
+#[cfg(feature = "debug_asset_server")]
+#[macro_export]
+macro_rules! load_internal_asset {
+    ($app: ident, $handle: ident, $path_str: expr, $loader: expr) => {{
+        {
+            let mut debug_app = $app
+                .world
+                .get_non_send_resource_mut::<bevy_asset::debug_asset_server::DebugAssetApp>()
+                .unwrap();
+            bevy_asset::debug_asset_server::register_handle_with_loader(
+                $loader,
+                &mut debug_app,
+                $handle,
+                file!(),
+                $path_str,
+            );
+        }
+        let mut assets = $app
+            .world
+            .get_resource_mut::<bevy_asset::Assets<_>>()
+            .unwrap();
+        assets.set_untracked($handle, ($loader)(include_str!($path_str)));
+    }};
+}
+
+#[cfg(not(feature = "debug_asset_server"))]
+#[macro_export]
+macro_rules! load_internal_asset {
+    ($app: ident, $handle: ident, $path_str: expr, $loader: expr) => {{
+        let mut assets = $app
+            .world
+            .get_resource_mut::<bevy_asset::Assets<_>>()
+            .unwrap();
+        assets.set_untracked($handle, ($loader)(include_str!($path_str)));
+    }};
+}
+
 #[cfg(test)]
 mod tests {
     use bevy_app::App;

crates/bevy_asset/src/debug_asset_server.rs

@@ -0,0 +1,138 @@
+use bevy_app::{App, Events, Plugin};
+use bevy_ecs::{
+    schedule::SystemLabel,
+    system::{NonSendMut, Res, ResMut, SystemState},
+};
+use bevy_tasks::{IoTaskPool, TaskPoolBuilder};
+use bevy_utils::HashMap;
+use std::{
+    ops::{Deref, DerefMut},
+    path::Path,
+};
+
+use crate::{
+    Asset, AssetEvent, AssetPlugin, AssetServer, AssetServerSettings, Assets, FileAssetIo, Handle,
+    HandleUntyped,
+};
+
+/// A "debug asset app", whose sole responsibility is hot reloading assets that are
+/// "internal" / compiled-in to Bevy Plugins.
+pub struct DebugAssetApp(App);
+
+impl Deref for DebugAssetApp {
+    type Target = App;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for DebugAssetApp {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+#[derive(SystemLabel, Debug, Clone, PartialEq, Eq, Hash)]
+pub struct DebugAssetAppRun;
+
+/// Facilitates the creation of a "debug asset app", whose sole responsibility is hot reloading
+/// assets that are "internal" / compiled-in to Bevy Plugins.
+/// Pair with [`load_internal_asset`](crate::load_internal_asset) to load "hot reloadable" assets
+/// The `debug_asset_server` feature flag must also be enabled for hot reloading to work.
+/// Currently only hot reloads assets stored in the `crates` folder.
+#[derive(Default)]
+pub struct DebugAssetServerPlugin;
+pub struct HandleMap<T: Asset> {
+    pub handles: HashMap<Handle<T>, Handle<T>>,
+}
+
+impl<T: Asset> Default for HandleMap<T> {
+    fn default() -> Self {
+        Self {
+            handles: Default::default(),
+        }
+    }
+}
+
+impl Plugin for DebugAssetServerPlugin {
+    fn build(&self, app: &mut bevy_app::App) {
+        let mut debug_asset_app = App::new();
+        debug_asset_app
+            .insert_resource(IoTaskPool(
+                TaskPoolBuilder::default()
+                    .num_threads(2)
+                    .thread_name("Debug Asset Server IO Task Pool".to_string())
+                    .build(),
+            ))
+            .insert_resource(AssetServerSettings {
+                asset_folder: "crates".to_string(),
+                watch_for_changes: true,
+            })
+            .add_plugin(AssetPlugin);
+        app.insert_non_send_resource(DebugAssetApp(debug_asset_app));
+        app.add_system(run_debug_asset_app);
+    }
+}
+
+fn run_debug_asset_app(mut debug_asset_app: NonSendMut<DebugAssetApp>) {
+    debug_asset_app.0.update();
+}
+
+pub(crate) fn sync_debug_assets<T: Asset + Clone>(
+    mut debug_asset_app: NonSendMut<DebugAssetApp>,
+    mut assets: ResMut<Assets<T>>,
+) {
+    let world = &mut debug_asset_app.0.world;
+    let mut state = SystemState::<(
+        Res<Events<AssetEvent<T>>>,
+        Res<HandleMap<T>>,
+        Res<Assets<T>>,
+    )>::new(world);
+    let (changed_shaders, handle_map, debug_assets) = state.get_mut(world);
+    for changed in changed_shaders.iter_current_update_events() {
+        let debug_handle = match changed {
+            AssetEvent::Created { handle } => handle,
+            AssetEvent::Modified { handle } => handle,
+            AssetEvent::Removed { .. } => continue,
+        };
+        if let Some(handle) = handle_map.handles.get(debug_handle) {
+            if let Some(debug_asset) = debug_assets.get(debug_handle) {
+                assets.set_untracked(handle, debug_asset.clone());
+            }
+        }
+    }
+}
+
+/// Uses the return type of the given loader to register the given handle with the appropriate type
+/// and load the asset with the given `path` and parent `file_path`.
+/// If this feels a bit odd ... thats because it is. This was built to improve the UX of the
+/// `load_internal_asset` macro.
+pub fn register_handle_with_loader<A: Asset>(
+    _loader: fn(&'static str) -> A,
+    app: &mut DebugAssetApp,
+    handle: HandleUntyped,
+    file_path: &str,
+    path: &'static str,
+) {
+    let mut state = SystemState::<(ResMut<HandleMap<A>>, Res<AssetServer>)>::new(&mut app.world);
+    let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
+    let manifest_dir_path = Path::new(&manifest_dir);
+    let (mut handle_map, asset_server) = state.get_mut(&mut app.world);
+    let asset_io = asset_server
+        .asset_io()
+        .downcast_ref::<FileAssetIo>()
+        .expect("The debug AssetServer only works with FileAssetIo-backed AssetServers");
+    let absolute_file_path = manifest_dir_path.join(
+        Path::new(file_path)
+            .parent()
+            .expect("file path must have a parent"),
+    );
+    let asset_folder_relative_path = absolute_file_path
+        .strip_prefix(asset_io.root_path())
+        .expect("The AssetIo root path should be a prefix of the absolute file path");
+    handle_map.handles.insert(
+        asset_server.load(asset_folder_relative_path.join(path)),
+        handle.clone_weak().typed::<A>(),
+    );
+}

crates/bevy_asset/src/io/file_asset_io.rs

@@ -62,6 +62,10 @@ impl FileAssetIo {
                 .unwrap()
         }
     }
+
+    pub fn root_path(&self) -> &PathBuf {
+        &self.root_path
+    }
 }
 
 impl AssetIo for FileAssetIo {

crates/bevy_asset/src/lib.rs

@@ -1,5 +1,7 @@
 mod asset_server;
 mod assets;
+#[cfg(feature = "debug_asset_server")]
+pub mod debug_asset_server;
 pub mod diagnostic;
 #[cfg(all(
     feature = "filesystem_watcher",

crates/bevy_internal/Cargo.toml

@@ -14,6 +14,7 @@ trace = [ "bevy_app/trace", "bevy_ecs/trace", "bevy_render/trace" ]
 trace_chrome = [ "bevy_log/tracing-chrome" ]
 trace_tracy = [ "bevy_log/tracing-tracy" ]
 wgpu_trace = ["bevy_render/wgpu_trace"]
+debug_asset_server = ["bevy_asset/debug_asset_server"]
 
 # Image format support for texture loading (PNG and HDR are enabled by default)
 hdr = ["bevy_render/hdr"]

crates/bevy_internal/src/default_plugins.rs

@@ -31,6 +31,8 @@ impl PluginGroup for DefaultPlugins {
         group.add(bevy_input::InputPlugin::default());
         group.add(bevy_window::WindowPlugin::default());
         group.add(bevy_asset::AssetPlugin::default());
+        #[cfg(feature = "debug_asset_server")]
+        group.add(bevy_asset::debug_asset_server::DebugAssetServerPlugin::default());
         group.add(bevy_scene::ScenePlugin::default());
 
         #[cfg(feature = "bevy_winit")]

crates/bevy_pbr/src/lib.rs

@@ -33,7 +33,7 @@ pub mod draw_3d_graph {
 }
 
 use bevy_app::prelude::*;
-use bevy_asset::{Assets, Handle, HandleUntyped};
+use bevy_asset::{load_internal_asset, Assets, Handle, HandleUntyped};
 use bevy_ecs::prelude::*;
 use bevy_reflect::TypeUuid;
 use bevy_render::{
@@ -57,14 +57,12 @@ pub struct PbrPlugin;
 
 impl Plugin for PbrPlugin {
     fn build(&self, app: &mut App) {
-        let mut shaders = app.world.get_resource_mut::<Assets<Shader>>().unwrap();
-        shaders.set_untracked(
-            PBR_SHADER_HANDLE,
-            Shader::from_wgsl(include_str!("render/pbr.wgsl")),
-        );
-        shaders.set_untracked(
+        load_internal_asset!(app, PBR_SHADER_HANDLE, "render/pbr.wgsl", Shader::from_wgsl);
+        load_internal_asset!(
+            app,
             SHADOW_SHADER_HANDLE,
-            Shader::from_wgsl(include_str!("render/depth.wgsl")),
+            "render/depth.wgsl",
+            Shader::from_wgsl
         );
 
         app.register_type::<CubemapVisibleEntities>()

crates/bevy_pbr/src/render/mesh.rs

@@ -3,7 +3,7 @@ use crate::{
     ViewClusterBindings, ViewLightsUniformOffset, ViewShadowBindings,
 };
 use bevy_app::Plugin;
-use bevy_asset::{Assets, Handle, HandleUntyped};
+use bevy_asset::{load_internal_asset, Handle, HandleUntyped};
 use bevy_ecs::{
     prelude::*,
     system::{lifetimeless::*, SystemParamItem},
@@ -35,18 +35,18 @@ pub const MESH_SHADER_HANDLE: HandleUntyped =
 
 impl Plugin for MeshRenderPlugin {
     fn build(&self, app: &mut bevy_app::App) {
-        let mut shaders = app.world.get_resource_mut::<Assets<Shader>>().unwrap();
-        shaders.set_untracked(
-            MESH_SHADER_HANDLE,
-            Shader::from_wgsl(include_str!("mesh.wgsl")),
-        );
-        shaders.set_untracked(
+        load_internal_asset!(app, MESH_SHADER_HANDLE, "mesh.wgsl", Shader::from_wgsl);
+        load_internal_asset!(
+            app,
             MESH_STRUCT_HANDLE,
-            Shader::from_wgsl(include_str!("mesh_struct.wgsl")),
+            "mesh_struct.wgsl",
+            Shader::from_wgsl
         );
-        shaders.set_untracked(
+        load_internal_asset!(
+            app,
             MESH_VIEW_BIND_GROUP_HANDLE,
-            Shader::from_wgsl(include_str!("mesh_view_bind_group.wgsl")),
+            "mesh_view_bind_group.wgsl",
+            Shader::from_wgsl
         );
 
         app.add_plugin(UniformComponentPlugin::<MeshUniform>::default());

crates/bevy_pbr/src/wireframe.rs

@@ -1,7 +1,7 @@
 use crate::MeshPipeline;
 use crate::{DrawMesh, MeshPipelineKey, MeshUniform, SetMeshBindGroup, SetMeshViewBindGroup};
 use bevy_app::Plugin;
-use bevy_asset::{Assets, Handle, HandleUntyped};
+use bevy_asset::{load_internal_asset, Handle, HandleUntyped};
 use bevy_core_pipeline::Opaque3d;
 use bevy_ecs::{prelude::*, reflect::ReflectComponent};
 use bevy_reflect::{Reflect, TypeUuid};
@@ -23,10 +23,11 @@ pub struct WireframePlugin;
 
 impl Plugin for WireframePlugin {
     fn build(&self, app: &mut bevy_app::App) {
-        let mut shaders = app.world.get_resource_mut::<Assets<Shader>>().unwrap();
-        shaders.set_untracked(
+        load_internal_asset!(
+            app,
             WIREFRAME_SHADER_HANDLE,
-            Shader::from_wgsl(include_str!("render/wireframe.wgsl")),
+            "render/wireframe.wgsl",
+            Shader::from_wgsl
         );
 
         app.init_resource::<WireframeConfig>();