1use fuchsia_inspect as finspect;
6use futures::FutureExt as _;
7use futures::future::BoxFuture;
8use std::collections::HashMap;
9use std::sync::{Arc, Weak};
10
11#[derive(Debug, Clone)]
43pub struct RootDirCache<S> {
44    non_meta_storage: S,
45    dirs: Arc<std::sync::Mutex<HashMap<fuchsia_hash::Hash, Weak<crate::RootDir<S>>>>>,
46}
47
48impl<S: crate::NonMetaStorage + Clone> RootDirCache<S> {
49    pub fn new(non_meta_storage: S) -> Self {
52        let dirs = Arc::new(std::sync::Mutex::new(HashMap::new()));
53        Self { non_meta_storage, dirs }
54    }
55
56    pub async fn get_or_insert(
63        &self,
64        hash: fuchsia_hash::Hash,
65        root_dir: Option<crate::RootDir<S>>,
66    ) -> Result<Arc<crate::RootDir<S>>, crate::Error> {
67        Ok(if let Some(root_dir) = self.get(&hash) {
68            root_dir
69        } else {
70            let dropper = Box::new(Dropper { dirs: Arc::downgrade(&self.dirs), hash });
72            let new_root_dir = match root_dir {
73                Some(mut root_dir) => match root_dir.set_dropper(dropper) {
74                    Ok(()) => Arc::new(root_dir),
75                    Err(_) => {
78                        return Err(crate::Error::DropperAlreadySet);
79                    }
80                },
81                None => {
82                    crate::RootDir::new_with_dropper(self.non_meta_storage.clone(), hash, dropper)
88                        .await?
89                }
90            };
91            use std::collections::hash_map::Entry::*;
92            let root_dir = match self.dirs.lock().expect("poisoned mutex").entry(hash) {
94                Occupied(mut o) => {
96                    let old_root_dir = o.get_mut();
97                    if let Some(old_root_dir) = old_root_dir.upgrade() {
98                        old_root_dir
99                    } else {
100                        *old_root_dir = Arc::downgrade(&new_root_dir);
101                        new_root_dir
102                    }
103                }
104                Vacant(v) => {
105                    v.insert(Arc::downgrade(&new_root_dir));
106                    new_root_dir
107                }
108            };
109            root_dir
110        })
111    }
112
113    pub fn get(&self, hash: &fuchsia_hash::Hash) -> Option<Arc<crate::RootDir<S>>> {
118        self.dirs.lock().expect("poisoned mutex").get(hash)?.upgrade()
119    }
120
121    pub fn list(&self) -> Vec<Arc<crate::RootDir<S>>> {
124        self.dirs.lock().expect("poisoned mutex").iter().filter_map(|(_, v)| v.upgrade()).collect()
125    }
126
127    pub fn record_lazy_inspect(
130        &self,
131    ) -> impl Fn() -> BoxFuture<'static, Result<finspect::Inspector, anyhow::Error>>
132    + Send
133    + Sync
134    + 'static {
135        let dirs = Arc::downgrade(&self.dirs);
136        move || {
137            let dirs = dirs.clone();
138            async move {
139                let inspector = finspect::Inspector::default();
140                if let Some(dirs) = dirs.upgrade() {
141                    let package_counts: HashMap<_, _> = {
142                        let dirs = dirs.lock().expect("poisoned mutex");
143                        dirs.iter().map(|(k, v)| (*k, v.strong_count() as u64)).collect()
144                    };
145                    let root = inspector.root();
146                    let () = package_counts.into_iter().for_each(|(pkg, count)| {
147                        root.record_child(pkg.to_string(), |n| n.record_uint("instances", count))
148                    });
149                }
150                Ok(inspector)
151            }
152            .boxed()
153        }
154    }
155}
156
157struct Dropper<S> {
159    dirs: Weak<std::sync::Mutex<HashMap<fuchsia_hash::Hash, Weak<crate::RootDir<S>>>>>,
160    hash: fuchsia_hash::Hash,
161}
162
163impl<S> Drop for Dropper<S> {
164    fn drop(&mut self) {
165        let Some(dirs) = self.dirs.upgrade() else {
166            return;
167        };
168        use std::collections::hash_map::Entry::*;
169        match dirs.lock().expect("poisoned mutex").entry(self.hash) {
170            Occupied(o) => {
171                if o.get().strong_count() == 0 {
173                    o.remove_entry();
174                }
175            }
176            Vacant(_) => (),
178        };
179    }
180}
181
182impl<S> std::fmt::Debug for Dropper<S> {
183    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
184        f.debug_struct("Dropper").field("dirs", &self.dirs).field("hash", &self.hash).finish()
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191    use assert_matches::assert_matches;
192    use diagnostics_assertions::assert_data_tree;
193    use fidl_fuchsia_io as fio;
194    use fuchsia_pkg_testing::PackageBuilder;
195    use fuchsia_pkg_testing::blobfs::Fake as FakeBlobfs;
196
197    #[fuchsia::test]
198    async fn get_or_insert_new_entry() {
199        let pkg = PackageBuilder::new("pkg-name").build().await.unwrap();
200        let (metafar_blob, _) = pkg.contents();
201        let (blobfs_fake, blobfs_client) = FakeBlobfs::new();
202        blobfs_fake.add_blob(metafar_blob.merkle, metafar_blob.contents);
203        let server = RootDirCache::new(blobfs_client);
204
205        let dir = server.get_or_insert(metafar_blob.merkle, None).await.unwrap();
206
207        assert_eq!(server.list().len(), 1);
208
209        drop(dir);
210        assert_eq!(server.list().len(), 0);
211        assert!(server.dirs.lock().expect("poisoned mutex").is_empty());
212    }
213
214    #[fuchsia::test]
215    async fn closing_package_connection_closes_package() {
216        let pkg = PackageBuilder::new("pkg-name").build().await.unwrap();
217        let (metafar_blob, _) = pkg.contents();
218        let (blobfs_fake, blobfs_client) = FakeBlobfs::new();
219        blobfs_fake.add_blob(metafar_blob.merkle, metafar_blob.contents);
220        let server = RootDirCache::new(blobfs_client);
221
222        let dir = server.get_or_insert(metafar_blob.merkle, None).await.unwrap();
223        let (proxy, server_end) = fidl::endpoints::create_proxy::<fio::DirectoryMarker>();
224        let scope = vfs::execution_scope::ExecutionScope::new();
225        vfs::directory::serve_on(dir, fio::PERM_READABLE, scope.clone(), server_end);
226        let _ = proxy
227            .get_attributes(Default::default())
228            .await
229            .expect("directory succesfully handling requests");
230        assert_eq!(server.list().len(), 1);
231
232        drop(proxy);
233        let () = scope.wait().await;
234        assert_eq!(server.list().len(), 0);
235        assert!(server.dirs.lock().expect("poisoned mutex").is_empty());
236    }
237
238    #[fuchsia::test]
239    async fn get_or_insert_existing_entry() {
240        let pkg = PackageBuilder::new("pkg-name").build().await.unwrap();
241        let (metafar_blob, _) = pkg.contents();
242        let (blobfs_fake, blobfs_client) = FakeBlobfs::new();
243        blobfs_fake.add_blob(metafar_blob.merkle, metafar_blob.contents);
244        let server = RootDirCache::new(blobfs_client);
245
246        let dir0 = server.get_or_insert(metafar_blob.merkle, None).await.unwrap();
247
248        let dir1 = server.get_or_insert(metafar_blob.merkle, None).await.unwrap();
249        assert_eq!(server.list().len(), 1);
250        assert_eq!(Arc::strong_count(&server.list()[0]), 3);
251
252        drop(dir0);
253        drop(dir1);
254        assert_eq!(server.list().len(), 0);
255        assert!(server.dirs.lock().expect("poisoned mutex").is_empty());
256    }
257
258    #[fuchsia::test]
259    async fn get_or_insert_provided_root_dir() {
260        let pkg = PackageBuilder::new("pkg-name").build().await.unwrap();
261        let (metafar_blob, _) = pkg.contents();
262        let (blobfs_fake, blobfs_client) = FakeBlobfs::new();
263        blobfs_fake.add_blob(metafar_blob.merkle, metafar_blob.contents);
264        let root_dir = crate::RootDir::new_raw(blobfs_client.clone(), metafar_blob.merkle, None)
265            .await
266            .unwrap();
267        blobfs_fake.delete_blob(metafar_blob.merkle);
268        let server = RootDirCache::new(blobfs_client);
269
270        let dir = server.get_or_insert(metafar_blob.merkle, Some(root_dir)).await.unwrap();
271        assert_eq!(server.list().len(), 1);
272
273        drop(dir);
274        assert_eq!(server.list().len(), 0);
275        assert!(server.dirs.lock().expect("poisoned mutex").is_empty());
276    }
277
278    #[fuchsia::test]
279    async fn get_or_insert_provided_root_dir_error_if_already_has_dropper() {
280        let pkg = PackageBuilder::new("pkg-name").build().await.unwrap();
281        let (metafar_blob, _) = pkg.contents();
282        let (blobfs_fake, blobfs_client) = FakeBlobfs::new();
283        blobfs_fake.add_blob(metafar_blob.merkle, metafar_blob.contents);
284        let root_dir =
285            crate::RootDir::new_raw(blobfs_client.clone(), metafar_blob.merkle, Some(Box::new(())))
286                .await
287                .unwrap();
288        let server = RootDirCache::new(blobfs_client);
289
290        assert_matches!(
291            server.get_or_insert(metafar_blob.merkle, Some(root_dir)).await,
292            Err(crate::Error::DropperAlreadySet)
293        );
294        assert!(server.dirs.lock().expect("poisoned mutex").is_empty());
295    }
296
297    #[fuchsia::test]
298    async fn get_or_insert_fails_if_root_dir_creation_fails() {
299        let (_blobfs_fake, blobfs_client) = FakeBlobfs::new();
300        let server = RootDirCache::new(blobfs_client);
301
302        assert_matches!(
303            server.get_or_insert([0; 32].into(), None).await,
304            Err(crate::Error::MissingMetaFar)
305        );
306        assert!(server.dirs.lock().expect("poisoned mutex").is_empty());
307    }
308
309    #[fuchsia::test]
310    async fn get_or_insert_concurrent_race_to_insert_new_root_dir() {
311        let pkg = PackageBuilder::new("pkg-name").build().await.unwrap();
312        let (metafar_blob, _) = pkg.contents();
313        let (blobfs_fake, blobfs_client) = FakeBlobfs::new();
314        blobfs_fake.add_blob(metafar_blob.merkle, metafar_blob.contents);
315        let server = RootDirCache::new(blobfs_client);
316
317        let fut0 = server.get_or_insert(metafar_blob.merkle, None);
318
319        let fut1 = server.get_or_insert(metafar_blob.merkle, None);
320
321        let (res0, res1) = futures::future::join(fut0, fut1).await;
322        let (dir0, dir1) = (res0.unwrap(), res1.unwrap());
323
324        assert_eq!(server.list().len(), 1);
325        assert_eq!(Arc::strong_count(&server.list()[0]), 3);
326
327        drop(dir0);
328        drop(dir1);
329        assert_eq!(server.list().len(), 0);
330        assert!(server.dirs.lock().expect("poisoned mutex").is_empty());
331    }
332
333    #[fuchsia::test]
334    async fn inspect() {
335        let pkg = PackageBuilder::new("pkg-name").build().await.unwrap();
336        let (metafar_blob, _) = pkg.contents();
337        let (blobfs_fake, blobfs_client) = FakeBlobfs::new();
338        blobfs_fake.add_blob(metafar_blob.merkle, metafar_blob.contents);
339        let server = RootDirCache::new(blobfs_client);
340        let _dir = server.get_or_insert(metafar_blob.merkle, None).await.unwrap();
341
342        let inspector = finspect::Inspector::default();
343        inspector.root().record_lazy_child("open-packages", server.record_lazy_inspect());
344
345        assert_data_tree!(inspector, root: {
346            "open-packages": {
347                pkg.hash().to_string() => {
348                    "instances": 1u64,
349                },
350            }
351        });
352    }
353}