Skip to main content

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