Store AssetSource as Arc in AssetSources.

andriyDev · andriyDev · commit 6463391c5aae · 2025-11-19T19:27:10.000-08:00
diff --git a/crates/bevy_asset/src/io/source.rs b/crates/bevy_asset/src/io/source.rs
@@ -165,13 +165,15 @@ impl AssetSourceBuilder {
         }
     }
 
-    /// Builds a new [`AssetSource`] with the given `id`. If `watch` is true, the unprocessed source will watch for changes.
-    /// If `watch_processed` is true, the processed source will watch for changes.
-    pub fn build(
+    /// Builds a new [`AssetSource`] with the given `id`. If `watch` is true, the unprocessed source
+    /// will watch for changes. If `watch_processed` is true, the processed source will watch for
+    /// changes. If `processing_state` is [`Some`], the processed reader will be gated on the state.
+    pub(crate) fn build(
         &mut self,
         id: AssetSourceId<'static>,
         watch: bool,
         watch_processed: bool,
+        processing_state: Option<Arc<ProcessingState>>,
     ) -> AssetSource {
         let reader = self.reader.as_mut()();
         let writer = self.writer.as_mut().and_then(|w| w(false));
@@ -222,6 +224,13 @@ impl AssetSourceBuilder {
                 }
             }
         }
+
+        if source.should_process()
+            && let Some(processing_state) = processing_state
+        {
+            source.gate_on_processor(processing_state);
+        }
+
         source
     }
 
@@ -355,25 +364,41 @@ impl AssetSourceBuilders {
         }
     }
 
-    /// Builds a new [`AssetSources`] collection. If `watch` is true, the unprocessed sources will watch for changes.
-    /// If `watch_processed` is true, the processed sources will watch for changes.
-    pub fn build_sources(&mut self, watch: bool, watch_processed: bool) -> AssetSources {
+    /// Builds a new [`AssetSources`] collection. If `watch` is true, the unprocessed sources will
+    /// watch for changes. If `watch_processed` is true, the processed sources will watch for
+    /// changes. If `processing_state` is [`Some`], the processed readers will be gated on the
+    /// processing state.
+    pub(crate) fn build_sources(
+        &mut self,
+        watch: bool,
+        watch_processed: bool,
+        processing_state: Option<Arc<ProcessingState>>,
+    ) -> AssetSources {
         let mut sources = <HashMap<_, _>>::default();
         for (id, source) in &mut self.sources {
             let source = source.build(
                 AssetSourceId::Name(id.clone_owned()),
                 watch,
                 watch_processed,
+                processing_state.clone(),
             );
-            sources.insert(id.clone_owned(), source);
+            sources.insert(id.clone_owned(), Arc::new(source));
         }
 
         AssetSources {
             sources,
             default: self
                 .default
                 .as_mut()
-                .map(|p| p.build(AssetSourceId::Default, watch, watch_processed))
+                .map(|p| {
+                    p.build(
+                        AssetSourceId::Default,
+                        watch,
+                        watch_processed,
+                        processing_state.clone(),
+                    )
+                })
+                .map(Arc::new)
                 .expect(MISSING_DEFAULT_SOURCE),
         }
     }
@@ -591,60 +616,43 @@ impl AssetSource {
 
 /// A collection of [`AssetSource`]s.
 pub struct AssetSources {
-    sources: HashMap<CowArc<'static, str>, AssetSource>,
-    default: AssetSource,
+    sources: HashMap<CowArc<'static, str>, Arc<AssetSource>>,
+    default: Arc<AssetSource>,
 }
 
 impl AssetSources {
     /// Gets the [`AssetSource`] with the given `id`, if it exists.
-    pub fn get<'a, 'b>(
-        &'a self,
-        id: impl Into<AssetSourceId<'b>>,
-    ) -> Result<&'a AssetSource, MissingAssetSourceError> {
+    pub fn get<'a>(
+        &self,
+        id: impl Into<AssetSourceId<'a>>,
+    ) -> Result<Arc<AssetSource>, MissingAssetSourceError> {
         match id.into().into_owned() {
-            AssetSourceId::Default => Ok(&self.default),
+            AssetSourceId::Default => Ok(self.default.clone()),
             AssetSourceId::Name(name) => self
                 .sources
                 .get(&name)
+                .cloned()
                 .ok_or(MissingAssetSourceError(AssetSourceId::Name(name))),
         }
     }
 
     /// Iterates all asset sources in the collection (including the default source).
-    pub fn iter(&self) -> impl Iterator<Item = &AssetSource> {
+    pub fn iter(&self) -> impl Iterator<Item = &Arc<AssetSource>> {
         self.sources.values().chain(Some(&self.default))
     }
 
-    /// Mutably iterates all asset sources in the collection (including the default source).
-    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut AssetSource> {
-        self.sources.values_mut().chain(Some(&mut self.default))
-    }
-
     /// Iterates all processed asset sources in the collection (including the default source).
-    pub fn iter_processed(&self) -> impl Iterator<Item = &AssetSource> {
+    pub fn iter_processed(&self) -> impl Iterator<Item = &Arc<AssetSource>> {
         self.iter().filter(|p| p.should_process())
     }
 
-    /// Mutably iterates all processed asset sources in the collection (including the default source).
-    pub fn iter_processed_mut(&mut self) -> impl Iterator<Item = &mut AssetSource> {
-        self.iter_mut().filter(|p| p.should_process())
-    }
-
     /// Iterates over the [`AssetSourceId`] of every [`AssetSource`] in the collection (including the default source).
     pub fn ids(&self) -> impl Iterator<Item = AssetSourceId<'static>> + '_ {
         self.sources
             .keys()
             .map(|k| AssetSourceId::Name(k.clone_owned()))
             .chain(Some(AssetSourceId::Default))
     }
-
-    /// This will cause processed [`AssetReader`](crate::io::AssetReader) futures (such as [`AssetReader::read`](crate::io::AssetReader::read)) to wait until
-    /// the [`AssetProcessor`](crate::AssetProcessor) has finished processing the requested asset.
-    pub(crate) fn gate_on_processor(&mut self, processing_state: Arc<ProcessingState>) {
-        for source in self.iter_processed_mut() {
-            source.gate_on_processor(processing_state.clone());
-        }
-    }
 }
 
 /// An error returned when an [`AssetSource`] does not exist for a given id.
diff --git a/crates/bevy_asset/src/lib.rs b/crates/bevy_asset/src/lib.rs
@@ -372,7 +372,7 @@ impl Plugin for AssetPlugin {
             match self.mode {
                 AssetMode::Unprocessed => {
                     let mut builders = app.world_mut().resource_mut::<AssetSourceBuilders>();
-                    let sources = builders.build_sources(watch, false);
+                    let sources = builders.build_sources(watch, false, None);
 
                     app.insert_resource(AssetServer::new_with_meta_check(
                         Arc::new(sources),
@@ -402,7 +402,7 @@ impl Plugin for AssetPlugin {
                         .add_systems(bevy_app::Startup, AssetProcessor::start);
                     } else {
                         let mut builders = app.world_mut().resource_mut::<AssetSourceBuilders>();
-                        let sources = builders.build_sources(false, watch);
+                        let sources = builders.build_sources(false, watch, None);
                         app.insert_resource(AssetServer::new_with_meta_check(
                             Arc::new(sources),
                             AssetServerMode::Processed,
diff --git a/crates/bevy_asset/src/loader_builders.rs b/crates/bevy_asset/src/loader_builders.rs
@@ -407,6 +407,7 @@ impl<'builder, 'reader, T> NestedLoader<'_, '_, T, Immediate<'builder, 'reader>>
             .write_infos()
             .stats
             .started_load_tasks += 1;
+        let source;
         let (mut meta, loader, mut reader) = if let Some(reader) = self.mode.reader {
             let loader = if let Some(asset_type_id) = asset_type_id {
                 self.load_context
@@ -430,10 +431,18 @@ impl<'builder, 'reader, T> NestedLoader<'_, '_, T, Immediate<'builder, 'reader>>
             let meta = loader.default_meta();
             (meta, loader, ReaderRef::Borrowed(reader))
         } else {
+            source = self
+                .load_context
+                .asset_server
+                .get_source(path.source())
+                .map_err(|err| LoadDirectError::LoadError {
+                    dependency: path.clone(),
+                    error: err.into(),
+                })?;
             let (meta, loader, reader) = self
                 .load_context
                 .asset_server
-                .get_meta_loader_and_reader(path, asset_type_id)
+                .get_meta_loader_and_reader(path, asset_type_id, &source)
                 .await
                 .map_err(|error| LoadDirectError::LoadError {
                     dependency: path.clone(),
diff --git a/crates/bevy_asset/src/processor/mod.rs b/crates/bevy_asset/src/processor/mod.rs
@@ -139,8 +139,7 @@ impl AssetProcessor {
         watch_processed: bool,
     ) -> (Self, Arc<AssetSources>) {
         let state = Arc::new(ProcessingState::new());
-        let mut sources = sources.build_sources(true, watch_processed);
-        sources.gate_on_processor(state.clone());
+        let sources = sources.build_sources(true, watch_processed, Some(state.clone()));
         let sources = Arc::new(sources);
 
         let data = Arc::new(AssetProcessorData::new(sources.clone(), state));
@@ -176,7 +175,7 @@ impl AssetProcessor {
     pub fn get_source<'a>(
         &self,
         id: impl Into<AssetSourceId<'a>>,
-    ) -> Result<&AssetSource, MissingAssetSourceError> {
+    ) -> Result<Arc<AssetSource>, MissingAssetSourceError> {
         self.data.sources.get(id.into())
     }
 
@@ -298,7 +297,7 @@ impl AssetProcessor {
                             return;
                         };
                         processor
-                            .handle_asset_source_event(source, event, &sender)
+                            .handle_asset_source_event(&source, event, &sender)
                             .await;
                     }
                 })
@@ -374,7 +373,9 @@ impl AssetProcessor {
                             let Ok(source) = processor.get_source(source_id) else {
                                 return;
                             };
-                            processor.process_asset(source, path, new_task_sender).await;
+                            processor
+                                .process_asset(&source, path, new_task_sender)
+                                .await;
                             // If the channel gets closed, that's ok. Just ignore it.
                             let _ = task_finished_sender.send(()).await;
                         })
diff --git a/crates/bevy_asset/src/server/mod.rs b/crates/bevy_asset/src/server/mod.rs
@@ -182,7 +182,7 @@ impl AssetServer {
     pub fn get_source<'a>(
         &self,
         source: impl Into<AssetSourceId<'a>>,
-    ) -> Result<&AssetSource, MissingAssetSourceError> {
+    ) -> Result<Arc<AssetSource>, MissingAssetSourceError> {
         self.data.sources.get(source.into())
     }
 
@@ -711,8 +711,19 @@ impl AssetServer {
 
         let path = path.into_owned();
         let path_clone = path.clone();
+        let source = self.get_source(path.source()).inspect_err(|e| {
+            // If there was an input handle, a "load" operation has already started, so we must
+            // produce a "failure" event, if we cannot find the source.
+            if let Some(handle) = &input_handle {
+                self.send_asset_event(InternalAssetEvent::Failed {
+                    index: handle.try_into().unwrap(),
+                    path: path.clone_owned(),
+                    error: e.clone().into(),
+                });
+            }
+        })?;
         let (mut meta, loader, mut reader) = self
-            .get_meta_loader_and_reader(&path_clone, input_handle_type_id)
+            .get_meta_loader_and_reader(&path_clone, input_handle_type_id, &source)
             .await
             .inspect_err(|e| {
                 // if there was an input handle, a "load" operation has already started, so we must produce a "failure" event, if
@@ -1418,6 +1429,7 @@ impl AssetServer {
         &'a self,
         asset_path: &'a AssetPath<'_>,
         asset_type_id: Option<TypeId>,
+        source: &'a AssetSource,
     ) -> Result<
         (
             Box<dyn AssetMetaDyn>,
@@ -1426,7 +1438,6 @@ impl AssetServer {
         ),
         AssetLoadError,
     > {
-        let source = self.get_source(asset_path.source())?;
         // NOTE: We grab the asset byte reader first to ensure this is transactional for AssetReaders like ProcessorGatedReader
         // The asset byte reader will "lock" the processed asset, preventing writes for the duration of the lock.
         // Then the meta reader, if meta exists, will correspond to the meta for the current "version" of the asset.
