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

elf_runner/
component_set.rs
