resource_accounting/
api.rs

1use std::{
2    collections::BTreeMap,
3    sync::{Arc, Mutex},
4};
5
6use saluki_api::{
7    extract::State,
8    response::IntoResponse,
9    routing::{get, Router},
10    APIHandler,
11};
12use serde::Serialize;
13
14use crate::{registry::ComponentMetadata, ResourceGroupRegistry, ResourceStatsSnapshot};
15
16#[derive(Serialize)]
17struct ComponentUsage {
18    minimum_required_bytes: usize,
19    firm_limit_bytes: usize,
20    actual_live_bytes: usize,
21}
22
23/// State used for the resource registry API handler.
24#[derive(Clone)]
25pub struct ResourceRegistryState {
26    inner: Arc<Mutex<ComponentMetadata>>,
27}
28
29impl ResourceRegistryState {
30    fn get_response(&self) -> String {
31        // The component registry is a nested structure, whereas the resource registry is a flat structure. We can
32        // only iterate via closure with the resource registry, so we do that, and for each entry, we look for it in
33        // the component registry. Luckily, the components in the resource registry use the same naming structure as
34        // the component registry supports for doing nested lookups in a single call, so we can just use that.
35
36        // TODO: This is only going to show components which have taken a token, which means some components in the
37        // registry, the ones whose bounds _are_ display when we verify bounds at startup, won't have an entry here
38        // because they don't use a token... _and_ even if we made sure to include all components in the registry, even
39        // if they didn't have an resource group associated, their usage might just be attributed to the root
40        // resource group so then you end up with components that say they have a minimum usage, etc, but show a
41        // current usage of zero bytes, etc.
42        //
43        // One option we could try is:
44        // - get the total minimum/firm bounds by rolling up all components
45        // - as we iterate over each resource group, we subtract the bounds of that component from the totals we
46        //   calculated before
47        // - at the end, we assign the remaining total bounds to the root component
48        //
49        // Essentially, because the memory usage for component with bounds that _don't_ use a token will inherently be
50        // attributed to the root resource group anyways, this would allow at least capturing those bounds for
51        // comparison against the otherwise-unattributed usage.
52        //
53        // Realistically, we _should_ still have our code set up to track all allocations for components that declare
54        // bounds, but this could be a reasonable stopgap.
55        let empty_snapshot = ResourceStatsSnapshot::empty();
56
57        let mut component_usage = BTreeMap::new();
58
59        let mut inner = self.inner.lock().unwrap();
60        ResourceGroupRegistry::global().visit_resource_groups(|component_name, component_stats| {
61            let component_meta = inner.get_or_create(component_name);
62            let component_meta = component_meta.lock().unwrap();
63            let bounds = component_meta.self_bounds();
64            let stats_snapshot = component_stats.snapshot_delta(&empty_snapshot);
65
66            component_usage.insert(
67                component_name.to_string(),
68                ComponentUsage {
69                    // TODO: figure out what's actually going on here, this might not be a valid change
70                    minimum_required_bytes: bounds.total_minimum_required_bytes(),
71                    firm_limit_bytes: bounds.total_firm_limit_bytes(),
72                    actual_live_bytes: stats_snapshot.live_bytes(),
73                },
74            );
75        });
76
77        serde_json::to_string(&component_usage).unwrap()
78    }
79}
80
81/// An API handler for reporting the resource usage and usage of all components.
82///
83/// This handler exposes a single route -- `/memory/status` -- which returns the overall bounds and live usage of each
84/// registered component.
85pub struct ResourceAPIHandler {
86    state: ResourceRegistryState,
87}
88
89impl ResourceAPIHandler {
90    pub(crate) fn from_state(state: Arc<Mutex<ComponentMetadata>>) -> Self {
91        Self {
92            state: ResourceRegistryState { inner: state },
93        }
94    }
95
96    async fn status_handler(State(state): State<ResourceRegistryState>) -> impl IntoResponse {
97        state.get_response()
98    }
99}
100
101impl APIHandler for ResourceAPIHandler {
102    type State = ResourceRegistryState;
103
104    fn generate_initial_state(&self) -> Self::State {
105        self.state.clone()
106    }
107
108    fn generate_routes(&self) -> Router<Self::State> {
109        Router::new().route("/memory/status", get(Self::status_handler))
110    }
111}
resource_accounting/api.rs

resource_accounting/
api.rs
