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

memory_accounting/
api.rs
