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

elf_runner/memory/
reporter.rs
