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

saluki_core/accounting/
api.rs
