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