khora_agents/asset_agent/
loader.rs

1// Copyright 2025 eraflo
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! A registry for asset loaders, enabling dynamic loading of different asset types by name.
16
17use anyhow::{anyhow, Result};
18use khora_core::asset::Asset;
19use khora_lanes::asset_lane::AssetLoaderLane;
20use khora_telemetry::{
21    metrics::registry::{CounterHandle, HistogramHandle},
22    MetricsRegistry, ScopedMetricTimer,
23};
24use std::{any::Any, collections::HashMap, sync::Arc};
25
26/// Internal trait for loading any asset type.
27trait AnyLoaderLane: Send + Sync {
28    fn load_any(&self, bytes: &[u8], metrics: &LoaderMetrics) -> Result<Box<dyn Any + Send>>;
29}
30
31/// A "wrapper" that takes a generic `AssetLoader<A>` and implements `AnyLoader`.
32struct AssetLoaderLaneWrapper<A: Asset, L: AssetLoaderLane<A>>(L, std::marker::PhantomData<A>);
33
34impl<A: Asset, L: AssetLoaderLane<A> + Send + Sync> AnyLoaderLane for AssetLoaderLaneWrapper<A, L> {
35    fn load_any(&self, bytes: &[u8], metrics: &LoaderMetrics) -> Result<Box<dyn Any + Send>> {
36        // Start the timer for this load operation.
37        let _timer = ScopedMetricTimer::new(&metrics.load_time_ms);
38
39        // Call the GENERIC and TYPE-SAFE load() method...
40        let asset: A = self.0.load(bytes).map_err(|e| anyhow!(e.to_string()))?;
41
42        // ...increment the asset loaded counter...
43        metrics.assets_loaded_total.increment()?;
44
45        // ...and return the result in a Box<dyn Any>.
46        Ok(Box::new(asset))
47    }
48}
49
50/// A collection of metric handles used by the loader registry.
51struct LoaderMetrics {
52    /// Histogram for tracking asset load times in milliseconds.
53    load_time_ms: HistogramHandle,
54    /// Counter for tracking the total number of assets loaded.
55    assets_loaded_total: CounterHandle,
56}
57
58impl LoaderMetrics {
59    fn new(registry: &MetricsRegistry) -> Self {
60        Self {
61            load_time_ms: registry
62                .register_histogram(
63                    "assets",
64                    "load_time",
65                    "Asset decoding time",
66                    "ms",
67                    vec![1.0, 5.0, 16.0, 33.0, 100.0, 500.0],
68                )
69                .expect("Failed to register asset load time metric"),
70            assets_loaded_total: registry
71                .register_counter(
72                    "assets",
73                    "loaded_total",
74                    "Total number of assets loaded from disk",
75                )
76                .expect("Failed to register asset count metric"),
77        }
78    }
79}
80
81/// The registry that manages complexity for the AssetAgent.
82pub(crate) struct AssetLoaderLaneRegistry {
83    /// Metrics handles for monitoring loader performance.
84    metrics: LoaderMetrics,
85    /// A map from asset type names to their corresponding loaders.
86    loaders: HashMap<String, Box<dyn AnyLoaderLane>>,
87}
88
89impl AssetLoaderLaneRegistry {
90    /// Creates a new `LoaderRegistry`.
91    pub(crate) fn new(metrics_registry: Arc<MetricsRegistry>) -> Self {
92        Self {
93            loaders: HashMap::new(),
94            metrics: LoaderMetrics::new(&metrics_registry),
95        }
96    }
97
98    /// Registers a new asset loader.
99    pub(crate) fn register<A: Asset>(
100        &mut self,
101        type_name: &str,
102        loader: impl AssetLoaderLane<A> + Send + Sync + 'static,
103    ) {
104        let wrapped = AssetLoaderLaneWrapper(loader, std::marker::PhantomData);
105        self.loaders
106            .insert(type_name.to_string(), Box::new(wrapped));
107    }
108
109    /// Loads an asset of the specified type from raw bytes.
110    pub(crate) fn load<A: Asset>(&self, type_name: &str, bytes: &[u8]) -> Result<A> {
111        let loader = self
112            .loaders
113            .get(type_name)
114            .ok_or_else(|| anyhow!("No loader registered for asset type '{}'", type_name))?;
115
116        let asset_any = loader.load_any(bytes, &self.metrics)?;
117
118        let asset_boxed = asset_any.downcast::<A>().map_err(|_| {
119            anyhow!(
120                "Loader for type '{}' returned a different asset type than requested.",
121                type_name
122            )
123        })?;
124
125        Ok(*asset_boxed)
126    }
127}
khora_agents/asset_agent/loader.rs

khora_agents/asset_agent/
loader.rs
