elf_runner/memory/
reporter.rs

1// Copyright 2023 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use attribution_server::{AttributionServer, AttributionServerHandle};
6use fidl::endpoints::{ControlHandle, DiscoverableProtocolMarker, RequestStream};
7use futures::TryStreamExt;
8use std::sync::Arc;
9use {fidl_fuchsia_io as fio, fidl_fuchsia_memory_attribution as fattribution};
10
11use crate::ComponentSet;
12use crate::component::ElfComponentInfo;
13
14pub struct MemoryReporter {
15    server: AttributionServerHandle,
16    components: Arc<ComponentSet>,
17}
18
19impl Drop for MemoryReporter {
20    fn drop(&mut self) {
21        self.components.set_callbacks(None, None);
22    }
23}
24
25impl MemoryReporter {
26    pub(crate) fn new(components: Arc<ComponentSet>) -> MemoryReporter {
27        let components_clone = components.clone();
28        let server = AttributionServer::new(Box::new(move || {
29            MemoryReporter::get_attribution(components_clone.as_ref())
30        }));
31        let new_component_publisher = server.new_publisher();
32        let deleted_component_publisher = server.new_publisher();
33        components.set_callbacks(
34            Some(Box::new(move |info| {
35                new_component_publisher.on_update(Self::build_new_attribution(info));
36            })),
37            Some(Box::new(move |token| {
38                deleted_component_publisher.on_update(vec![
39                    fattribution::AttributionUpdate::Remove(token.koid().unwrap().raw_koid()),
40                ]);
41            })),
42        );
43        MemoryReporter { server, components }
44    }
45
46    fn get_attribution(components: &ComponentSet) -> Vec<fattribution::AttributionUpdate> {
47        let mut attributions: Vec<fattribution::AttributionUpdate> = vec![];
48        components.visit(|component: &ElfComponentInfo, _id| {
49            let mut component_attributions = Self::build_new_attribution(component);
50            attributions.append(&mut component_attributions);
51        });
52        attributions
53    }
54
55    pub fn serve(&self, mut stream: fattribution::ProviderRequestStream) {
56        let subscriber = self.server.new_observer(stream.control_handle());
57        self.components.scope().spawn(async move {
58            while let Ok(Some(request)) = stream.try_next().await {
59                match request {
60                    fattribution::ProviderRequest::Get { responder } => {
61                        subscriber.next(responder);
62                    }
63                    fattribution::ProviderRequest::_UnknownMethod {
64                        ordinal,
65                        control_handle,
66                        ..
67                    } => {
68                        log::error!("Invalid request to AttributionProvider: {ordinal}");
69                        control_handle.shutdown_with_epitaph(zx::Status::INVALID_ARGS);
70                    }
71                }
72            }
73        });
74    }
75
76    fn build_new_attribution(component: &ElfComponentInfo) -> Vec<fattribution::AttributionUpdate> {
77        let instance_token = component.copy_instance_token().unwrap();
78        let instance_token_koid = instance_token.koid().unwrap().raw_koid();
79        let new_principal = fattribution::NewPrincipal {
80            identifier: Some(instance_token_koid),
81            description: Some(fattribution::Description::Component(instance_token)),
82            principal_type: Some(fattribution::PrincipalType::Runnable),
83            detailed_attribution: component.get_outgoing_directory().and_then(
84                |outgoing_directory| {
85                    let (server, client) = fidl::Channel::create();
86                    fdio::open_at(
87                        outgoing_directory.channel(),
88                        &format!("svc/{}", fattribution::ProviderMarker::PROTOCOL_NAME),
89                        fio::Flags::empty(),
90                        server,
91                    )
92                    .unwrap();
93                    let provider =
94                        fidl::endpoints::ClientEnd::<fattribution::ProviderMarker>::new(client);
95                    Some(provider)
96                },
97            ),
98            ..Default::default()
99        };
100        let attribution = fattribution::UpdatedPrincipal {
101            identifier: Some(instance_token_koid),
102            resources: Some(fattribution::Resources::Data(fattribution::Data {
103                resources: vec![fattribution::Resource::KernelObject(
104                    component.copy_job().proc().koid().unwrap().raw_koid(),
105                )],
106            })),
107            ..Default::default()
108        };
109        vec![
110            fattribution::AttributionUpdate::Add(new_principal),
111            fattribution::AttributionUpdate::Update(attribution),
112        ]
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119    use crate::tests::{lifecycle_startinfo, new_elf_runner_for_test};
120    use cm_config::SecurityPolicy;
121    use futures::FutureExt;
122    use moniker::Moniker;
123    use routing::policy::ScopedPolicyChecker;
124    use {
125        fidl_fuchsia_component_runner as fcrunner, fidl_fuchsia_io as fio, fuchsia_async as fasync,
126    };
127
128    /// Test that the ELF runner can tell us about the resources used by the component it runs.
129    #[test]
130    fn test_attribute_memory() {
131        // TODO(https://fxbug.dev/463680736) remove once deadlocks addressed
132        fuchsia_sync::suppress_lock_cycle_panics();
133
134        let mut exec = fasync::TestExecutor::new();
135        let (_runtime_dir, runtime_dir_server) =
136            fidl::endpoints::create_proxy::<fio::DirectoryMarker>();
137        let start_info = lifecycle_startinfo(runtime_dir_server);
138
139        let runner = new_elf_runner_for_test();
140        let (snapshot_provider, snapshot_request_stream) =
141            fidl::endpoints::create_proxy_and_stream::<fattribution::ProviderMarker>();
142        runner.serve_memory_reporter(snapshot_request_stream);
143
144        // Run a component.
145        let runner = runner.get_scoped_runner(ScopedPolicyChecker::new(
146            Arc::new(SecurityPolicy::default()),
147            Moniker::try_from("foo/bar").unwrap(),
148        ));
149        let (_controller, server_controller) =
150            fidl::endpoints::create_proxy::<fcrunner::ComponentControllerMarker>();
151        exec.run_singlethreaded(&mut runner.start(start_info, server_controller).boxed());
152
153        // Ask about the memory usage of components.
154        let attributions =
155            exec.run_singlethreaded(snapshot_provider.get()).unwrap().unwrap().attributions;
156        assert!(attributions.is_some());
157
158        let attributions_vec = attributions.unwrap();
159        // It should contain one component, the one we just launched.
160        assert_eq!(attributions_vec.len(), 2);
161        let new_attrib = attributions_vec.get(0).unwrap();
162        let fattribution::AttributionUpdate::Add(added_principal) = new_attrib else {
163            panic!("Not a new principal");
164        };
165        assert_eq!(added_principal.principal_type, Some(fattribution::PrincipalType::Runnable));
166
167        // Its resource is a single job.
168        let update_attrib = attributions_vec.get(1).unwrap();
169        let fattribution::AttributionUpdate::Update(_) = update_attrib else {
170            panic!("Not an update");
171        };
172    }
173}