use crate::assets::ImageAssetData;
use crate::assets::{BufferAsset, ImageAsset, MaterialAsset};
use crate::{
AssetLookup, AssetTypeHandler, BufferAssetData, MaterialInstanceSlotAssignment,
RafxGenericLoadEventHandler,
};
use hydrate_base::handle::{ArtifactHandle, Handle, LoadState};
use rafx_framework::{
DescriptorSetAllocatorMetrics, DescriptorSetAllocatorProvider, DescriptorSetAllocatorRef,
DescriptorSetLayoutResource, DescriptorSetWriteSet, DynResourceAllocatorSet,
GraphicsPipelineCache, MaterialPass, RenderResources, ResourceArc, SlotNameLookup,
};
use crate::assets::buffer::BufferAssetTypeHandler;
use crate::assets::compute_pipeline::ComputePipelineAssetTypeHandler;
use crate::assets::graphics_pipeline::{
MaterialAssetTypeHandler, MaterialInstanceAssetTypeHandler, SamplerAssetTypeHandler,
};
use crate::assets::image::ImageAssetTypeHandler;
use crate::assets::shader::ShaderAssetTypeHandler;
use crate::hydrate_impl::AssetResource;
use fnv::FnvHashMap;
use rafx_api::{RafxDeviceContext, RafxQueue, RafxResult};
use rafx_framework::descriptor_sets::{
DescriptorSetElementKey, DescriptorSetWriteElementBuffer, DescriptorSetWriteElementBufferData,
DescriptorSetWriteElementImage,
};
use rafx_framework::render_features::RenderRegistry;
use rafx_framework::upload::{UploadQueue, UploadQueueConfig, UploadQueueContext};
use rafx_framework::DescriptorSetAllocator;
use rafx_framework::DynCommandPoolAllocator;
use rafx_framework::DynResourceAllocatorSetProvider;
use rafx_framework::ResourceLookupSet;
use rafx_framework::{ResourceManager, ResourceManagerMetrics};
use std::any::TypeId;
use std::sync::Arc;
#[derive(Debug)]
pub struct AssetManagerMetrics {
pub resource_manager_metrics: ResourceManagerMetrics,
pub material_instance_descriptor_sets_metrics: DescriptorSetAllocatorMetrics,
}
pub struct AssetManagerLoaders {
pub image_loader: RafxGenericLoadEventHandler<ImageAssetData, ImageAsset>,
pub buffer_loader: RafxGenericLoadEventHandler<BufferAssetData, BufferAsset>,
}
pub struct AssetManager {
device_context: RafxDeviceContext,
resource_manager: ResourceManager,
upload_queue: UploadQueue,
material_instance_descriptor_sets: DescriptorSetAllocator,
graphics_queue: RafxQueue,
transfer_queue: RafxQueue,
asset_types: FnvHashMap<TypeId, Box<dyn AssetTypeHandler>>,
asset_registration_order: Arc<Vec<TypeId>>,
}
impl AssetManager {
pub fn new(
device_context: &RafxDeviceContext,
render_registry: &RenderRegistry,
upload_queue_config: UploadQueueConfig,
graphics_queue: &RafxQueue,
transfer_queue: &RafxQueue,
) -> RafxResult<Self> {
let resource_manager = ResourceManager::new(device_context, render_registry);
let upload_queue = UploadQueue::new(
device_context,
upload_queue_config,
graphics_queue.clone(),
transfer_queue.clone(),
)?;
Ok(AssetManager {
device_context: device_context.clone(),
resource_manager,
upload_queue,
material_instance_descriptor_sets: DescriptorSetAllocator::new(device_context),
graphics_queue: graphics_queue.clone(),
transfer_queue: transfer_queue.clone(),
asset_types: Default::default(),
asset_registration_order: Default::default(),
})
}
pub fn register_asset_type(
&mut self,
asset_type: Box<dyn AssetTypeHandler>,
) -> RafxResult<()> {
let mut asset_registration_order = (*self.asset_registration_order).clone();
asset_registration_order.push(asset_type.asset_type_id());
self.asset_registration_order = Arc::new(asset_registration_order);
let old = self
.asset_types
.insert(asset_type.asset_type_id(), asset_type);
assert!(old.is_none());
Ok(())
}
pub fn register_default_asset_types(
&mut self,
asset_resource: &mut AssetResource,
_render_resources: &mut RenderResources,
) -> RafxResult<()> {
let asset_type = ShaderAssetTypeHandler::create(self, asset_resource)?;
self.register_asset_type(asset_type)?;
let asset_type = ComputePipelineAssetTypeHandler::create(self, asset_resource)?;
self.register_asset_type(asset_type)?;
let asset_type = MaterialAssetTypeHandler::create(self, asset_resource)?;
self.register_asset_type(asset_type)?;
let asset_type = MaterialInstanceAssetTypeHandler::create(self, asset_resource)?;
self.register_asset_type(asset_type)?;
let asset_type = SamplerAssetTypeHandler::create(self, asset_resource)?;
self.register_asset_type(asset_type)?;
let asset_type = ImageAssetTypeHandler::create(self, asset_resource)?;
self.register_asset_type(asset_type)?;
let asset_type = BufferAssetTypeHandler::create(self, asset_resource)?;
self.register_asset_type(asset_type)?;
Ok(())
}
pub fn committed_asset<AssetT: 'static>(
&self,
handle: &Handle<AssetT>,
) -> Option<&AssetT> {
let asset_type = self.asset_types.get(&TypeId::of::<AssetT>())?;
asset_type
.asset_lookup()
.downcast_ref::<AssetLookup<AssetT>>()
.unwrap()
.get_committed(handle.resolved_load_handle())
}
pub fn latest_asset<AssetT: 'static>(
&self,
handle: &Handle<AssetT>,
) -> Option<&AssetT> {
let asset_type = self.asset_types.get(&TypeId::of::<AssetT>())?;
asset_type
.asset_lookup()
.downcast_ref::<AssetLookup<AssetT>>()
.unwrap()
.get_latest(handle.resolved_load_handle())
}
#[profiling::function]
pub fn wait_for_asset_to_load<
T,
TickFn: FnMut(&mut AssetManager, &mut AssetResource) -> RafxResult<()>,
>(
&mut self,
asset_handle: &Handle<T>,
asset_resource: &mut AssetResource,
asset_name: &str,
mut tick_fn: TickFn,
) -> RafxResult<()> {
const PRINT_INTERVAL: std::time::Duration = std::time::Duration::from_millis(1000);
let mut last_print_time: Option<rafx_base::Instant> = None;
fn on_interval<F: Fn()>(
interval: std::time::Duration,
last_time: &mut Option<rafx_base::Instant>,
f: F,
) {
let now = rafx_base::Instant::now();
if last_time.is_none() || now - last_time.unwrap() >= interval {
(f)();
*last_time = Some(now);
}
}
loop {
asset_resource.update();
self.update_asset_loaders()?;
(tick_fn)(self, asset_resource)?;
match asset_handle.load_state(asset_resource.loader()) {
LoadState::Loaded => {
break Ok(());
}
state @ _ => {
on_interval(PRINT_INTERVAL, &mut last_print_time, || {
let artifact_id = asset_handle.artifact_id(asset_resource.loader());
log::info!(
"blocked waiting for asset to resolve Name={} Handle={:?} ArtifactId={:?} State={:?}",
asset_name,
asset_handle.resolved_load_handle().id,
artifact_id,
state,
);
});
}
}
}
}
pub fn device_context(&self) -> &RafxDeviceContext {
&self.device_context
}
pub fn graphics_queue(&self) -> &RafxQueue {
&self.graphics_queue
}
pub fn transfer_queue(&self) -> &RafxQueue {
&self.transfer_queue
}
pub fn resource_manager(&self) -> &ResourceManager {
&self.resource_manager
}
pub fn resource_manager_mut(&mut self) -> &mut ResourceManager {
&mut self.resource_manager
}
pub fn resources(&self) -> &ResourceLookupSet {
self.resource_manager.resources()
}
pub fn graphics_pipeline_cache(&self) -> &GraphicsPipelineCache {
self.resource_manager.graphics_pipeline_cache()
}
pub fn dyn_command_pool_allocator(&self) -> &DynCommandPoolAllocator {
self.resource_manager.dyn_command_pool_allocator()
}
pub fn create_dyn_resource_allocator_set(&self) -> DynResourceAllocatorSet {
self.resource_manager.create_dyn_resource_allocator_set()
}
pub fn create_dyn_resource_allocator_provider(&self) -> DynResourceAllocatorSetProvider {
self.resource_manager
.create_dyn_resource_allocator_provider()
}
pub fn create_descriptor_set_allocator(&self) -> DescriptorSetAllocatorRef {
self.resource_manager.create_descriptor_set_allocator()
}
pub fn create_descriptor_set_allocator_provider(&self) -> DescriptorSetAllocatorProvider {
self.resource_manager
.create_descriptor_set_allocator_provider()
}
pub(crate) fn material_instance_descriptor_sets_mut(&mut self) -> &mut DescriptorSetAllocator {
&mut self.material_instance_descriptor_sets
}
pub fn upload_queue_context(&self) -> UploadQueueContext {
self.upload_queue.upload_queue_context()
}
pub fn get_descriptor_set_layout_for_pass(
&self,
handle: &Handle<MaterialAsset>,
pass_index: usize,
layout_index: usize,
) -> Option<ResourceArc<DescriptorSetLayoutResource>> {
self.committed_asset(&handle)
.and_then(|x| x.passes.get(pass_index))
.map(|x| {
x.material_pass_resource.get_raw().descriptor_set_layouts[layout_index].clone()
})
}
#[profiling::function]
pub fn update_asset_loaders(&mut self) -> RafxResult<()> {
for asset_type in &*self.asset_registration_order.clone() {
let mut asset_type = self.asset_types.remove(asset_type).unwrap();
asset_type.process_load_requests(self)?;
self.asset_types
.insert(asset_type.asset_type_id(), asset_type);
}
self.upload_queue.update()?;
Ok(())
}
pub fn on_begin_frame(&mut self) -> RafxResult<()> {
self.material_instance_descriptor_sets.flush_changes()
}
#[profiling::function]
pub fn on_frame_complete(&mut self) -> RafxResult<()> {
for (_, asset_type) in &mut self.asset_types {
asset_type.on_frame_complete()?;
}
self.resource_manager.on_frame_complete()?;
self.material_instance_descriptor_sets.on_frame_complete();
Ok(())
}
pub fn metrics(&self) -> AssetManagerMetrics {
let resource_manager_metrics = self.resource_manager.metrics();
let material_instance_descriptor_sets_metrics =
self.material_instance_descriptor_sets.metrics();
AssetManagerMetrics {
resource_manager_metrics,
material_instance_descriptor_sets_metrics,
}
}
#[profiling::function]
pub fn apply_material_instance_slot_assignment(
&self,
slot_assignment: &MaterialInstanceSlotAssignment,
pass_slot_name_lookup: &SlotNameLookup,
resources: &ResourceLookupSet,
material_pass_write_set: &mut Vec<DescriptorSetWriteSet>,
) -> RafxResult<()> {
if let Some(slot_locations) = pass_slot_name_lookup.get(&slot_assignment.slot_name) {
for location in slot_locations {
log::trace!(
"Apply write to location {:?} via slot {}",
location,
slot_assignment.slot_name
);
let layout_descriptor_set_writes =
&mut material_pass_write_set[location.layout_index as usize];
let write = layout_descriptor_set_writes
.elements
.get_mut(&DescriptorSetElementKey {
dst_binding: location.binding_index,
array_index: slot_assignment.array_index,
})
.unwrap();
let what_to_bind = rafx_framework::descriptor_sets::what_to_bind(write);
if what_to_bind.bind_images || what_to_bind.bind_samplers {
let mut write_image = DescriptorSetWriteElementImage {
image_view: None,
sampler: None,
};
if what_to_bind.bind_images {
if let Some(image) = &slot_assignment.image {
let loaded_image = self.latest_asset(&image).unwrap();
write_image.image_view =
Some(rafx_framework::descriptor_sets::DescriptorSetWriteElementImageValue::Resource(
loaded_image.image_view.clone(),
));
}
}
if what_to_bind.bind_samplers {
if let Some(sampler) = &slot_assignment.sampler {
let sampler = resources.get_or_create_sampler(sampler)?;
write_image.sampler = Some(sampler);
}
}
write.image_info = write_image;
}
if what_to_bind.bind_buffers {
let mut write_buffer = DescriptorSetWriteElementBuffer { buffer: None };
if let Some(buffer_data) = &slot_assignment.buffer_data {
write_buffer.buffer = Some(DescriptorSetWriteElementBufferData::Data(
buffer_data.clone(),
));
}
write.buffer_info = write_buffer;
}
}
}
Ok(())
}
pub fn create_write_sets_for_material_instance_pass(
&self,
pass: &MaterialPass,
slots: &[MaterialInstanceSlotAssignment],
resources: &ResourceLookupSet,
) -> RafxResult<Vec<DescriptorSetWriteSet>> {
let mut pass_descriptor_set_writes =
pass.create_uninitialized_write_sets_for_material_pass();
for slot in slots {
self.apply_material_instance_slot_assignment(
slot,
&pass.pass_slot_name_lookup,
resources,
&mut pass_descriptor_set_writes,
)?;
}
Ok(pass_descriptor_set_writes)
}
}
impl Drop for AssetManager {
fn drop(&mut self) {
log::info!("Cleaning up asset manager");
log::trace!("Asset Manager Metrics:\n{:#?}", self.metrics());
self.transfer_queue.wait_for_queue_idle().unwrap();
self.graphics_queue.wait_for_queue_idle().unwrap();
for asset_type in self.asset_registration_order.iter().rev() {
self.asset_types.remove(asset_type).unwrap();
}
self.material_instance_descriptor_sets.destroy().unwrap();
log::info!("Dropping asset manager");
log::trace!("Asset Manager Metrics:\n{:#?}", self.metrics());
}
}