elf_runner/
component_set.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 crate::{ElfComponent, ElfComponentInfo};
6use fuchsia_async as fasync;
7use fuchsia_sync::Mutex;
8use id::Id;
9use std::collections::HashMap;
10use std::sync::{Arc, Weak};
11
12/// [`ComponentSet`] tracks all the components executing inside an ELF runner,
13/// and presents an iterator over those components. It does this under the
14/// constraint that each component may go out of scope concurrently due to
15/// stopping on its own, or being stopped by the `ComponentController` protocol
16/// or any other reason.
17#[derive(Default)]
18pub struct ComponentSet {
19    inner: Mutex<HashMap<Id, Weak<ElfComponentInfo>>>,
20
21    /// Callbacks for component insertion and removal.
22    ///
23    /// These callbacks need to be behind a [Mutex] as it isn't guaranteed they
24    /// will be called on the same thread as [ComponentSet].
25    callbacks: Mutex<ComponentSetCallbacks>,
26}
27
28#[derive(Default)]
29struct ComponentSetCallbacks {
30    on_new_component: Option<Box<dyn Fn(&ElfComponentInfo) -> () + Send>>,
31    on_removed_component: Option<Box<dyn Fn(zx::Event) -> () + Send>>,
32}
33
34impl ComponentSet {
35    pub fn new() -> Arc<Self> {
36        Default::default()
37    }
38
39    pub fn set_callbacks(
40        &self,
41        on_new_component: Option<Box<dyn Fn(&ElfComponentInfo) -> () + Send>>,
42        on_removed_component: Option<Box<dyn Fn(zx::Event) -> () + Send>>,
43    ) {
44        let mut locked_callbacks = self.callbacks.lock();
45        locked_callbacks.on_new_component = on_new_component;
46        locked_callbacks.on_removed_component = on_removed_component;
47    }
48
49    /// Adds a component to the set.
50    ///
51    /// The component will remove itself from the set when it is dropped.
52    pub fn add(self: Arc<Self>, component: &mut ElfComponent) {
53        let mut inner = self.inner.lock();
54        let component_set = Arc::downgrade(&self.clone());
55        let id = Id::new(component.info().get_moniker().clone());
56        let id_clone = id.clone();
57        component.set_on_drop(ComponentSet::on_component_drop(component_set, id_clone));
58        inner.insert(id, Arc::downgrade(component.info()));
59        if let Some(cb) = self.callbacks.lock().on_new_component.as_ref() {
60            cb(component.info().as_ref());
61        }
62    }
63
64    /// Invokes `visitor` over all [`ElfComponentInfo`] objects corresponding to
65    /// components that are currently running. Note that this is fundamentally racy
66    /// as a component could be stopping imminently during or after the visit.
67    pub fn visit(&self, mut visitor: impl FnMut(&ElfComponentInfo, Id)) {
68        // Cloning the map helps us avoid holding the lock when iterating over it, in
69        // particular while calling the callback.
70        let components = &self.inner.lock().clone();
71        for (id, component) in components.iter() {
72            let Some(component) = component.upgrade() else {
73                continue;
74            };
75            visitor(&component, id.clone());
76        }
77    }
78
79    /// Callback when a component is removed from the set.
80    fn on_component_drop(
81        component_set: Weak<ComponentSet>,
82        id_clone: Id,
83    ) -> impl FnOnce(&ElfComponentInfo) {
84        move |info| {
85            let token = info.copy_instance_token().unwrap();
86            fasync::Task::spawn(async move {
87                let Some(component_set) = component_set.upgrade() else {
88                    return;
89                };
90                {
91                    let mut locked_inner = component_set.inner.lock();
92                    locked_inner.remove(&id_clone);
93                }
94                if let Some(cb) = component_set.callbacks.lock().on_removed_component.as_ref() {
95                    cb(token);
96                };
97            })
98            .detach()
99        }
100    }
101}
102
103/// An identifier for running components that is unique within an ELF runner.
104pub mod id {
105    use moniker::Moniker;
106    use std::fmt::Display;
107    use std::sync::atomic::{AtomicU64, Ordering};
108
109    /// TODO(https://fxbug.dev/316036032): store a moniker token instead, to eliminate
110    /// ELF runner's visibility into component monikers.
111    #[derive(Eq, Hash, PartialEq, Clone, Debug)]
112    pub struct Id(Moniker, u64);
113
114    static NEXT_ID: AtomicU64 = AtomicU64::new(0);
115
116    impl Id {
117        pub fn new(moniker: Moniker) -> Id {
118            Id(moniker, NEXT_ID.fetch_add(1, Ordering::SeqCst))
119        }
120    }
121
122    impl Display for Id {
123        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124            f.write_fmt(format_args!("{}, {}", self.0, self.1))
125        }
126    }
127
128    impl From<Id> for String {
129        fn from(value: Id) -> Self {
130            format!("{value}")
131        }
132    }
133
134    impl TryFrom<String> for Id {
135        type Error = anyhow::Error;
136
137        fn try_from(value: String) -> Result<Self, Self::Error> {
138            let Some((moniker, counter)) = value.split_once(", ") else {
139                anyhow::bail!("Expected comma separated string, got {value}");
140            };
141            let moniker: Moniker = moniker.parse()?;
142            let counter: u64 = counter.parse()?;
143            Ok(Self(moniker, counter))
144        }
145    }
146
147    #[cfg(test)]
148    mod tests {
149        use super::*;
150
151        #[test]
152        fn test_get_id() {
153            let id1 = Id::new(Moniker::try_from("foo/bar").unwrap());
154            let id2 = Id::new(Moniker::try_from("foo/bar").unwrap());
155            assert_ne!(id1, id2);
156        }
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    use futures::FutureExt;
165    use moniker::Moniker;
166    use std::future;
167    use std::sync::atomic::{AtomicUsize, Ordering};
168    use std::task::Poll;
169
170    use crate::runtime_dir::RuntimeDirectory;
171    use crate::Job;
172
173    #[test]
174    fn test_add_remove_component() {
175        // Use a test executor so that we can run until stalled.
176        let mut exec = fasync::TestExecutor::new();
177        let components = ComponentSet::new();
178
179        // The component set starts out empty.
180        let count = Arc::new(AtomicUsize::new(0));
181        let components_clone = components.clone();
182        components_clone.visit(|_, _| {
183            count.fetch_add(1, Ordering::SeqCst);
184        });
185        assert_eq!(count.load(Ordering::SeqCst), 0);
186
187        // After adding, it should contain one component.
188        let mut fake_component = make_fake_component();
189        components.clone().add(&mut fake_component);
190
191        let count = Arc::new(AtomicUsize::new(0));
192        let components_clone = components.clone();
193        components_clone.visit(|_, _| {
194            count.fetch_add(1, Ordering::SeqCst);
195        });
196        assert_eq!(count.load(Ordering::SeqCst), 1);
197
198        // After dropping that component, it should eventually contain zero components.
199        drop(fake_component);
200        let mut fut = async {
201            let _: Poll<()> =
202                fasync::TestExecutor::poll_until_stalled(future::pending::<()>()).await;
203        }
204        .boxed();
205        assert!(exec.run_until_stalled(&mut fut).is_ready());
206
207        let count = Arc::new(AtomicUsize::new(0));
208        let components_clone = components.clone();
209        components_clone.visit(|_, _| {
210            count.fetch_add(1, Ordering::SeqCst);
211        });
212        assert_eq!(count.load(Ordering::SeqCst), 0);
213    }
214
215    fn make_fake_component() -> ElfComponent {
216        let runtime_dir = RuntimeDirectory::empty();
217        let job = Job::Single(fuchsia_runtime::job_default().create_child_job().unwrap());
218        let process = fuchsia_runtime::process_self().duplicate(zx::Rights::SAME_RIGHTS).unwrap();
219        let lifecycle_channel = None;
220        let main_process_critical = false;
221        let tasks = vec![];
222        let component_url = "hello".to_string();
223        let fake_component = ElfComponent::new(
224            runtime_dir,
225            Moniker::default(),
226            job,
227            process,
228            lifecycle_channel,
229            main_process_critical,
230            tasks,
231            component_url,
232            None,
233            Default::default(),
234            zx::Event::create(),
235        );
236        fake_component
237    }
238}
elf_runner/component_set.rs

elf_runner/
component_set.rs
