block_server/
bin.rs

1// Copyright 2025 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 slab::Slab;
6
7// Bin is a epoch-like data structure for dropping things once all previous references have been
8// released.  For the C interface we hand out unowned VMOs whose lifetime is tied to that of the
9// request.  Rather than retain all the VMOs that might be used by a request, we just hold on to
10// closed VMOs until *all* preceding requests have finished.
11pub struct Bin<T> {
12    bin: Slab<BinEntry<T>>,
13    current_key: usize,
14}
15
16struct BinEntry<T> {
17    ref_count: usize,
18    data_and_next: Option<(T, usize)>,
19}
20
21impl<T> Bin<T> {
22    pub fn new() -> Self {
23        Self {
24            bin: Slab::from_iter([(0, BinEntry { ref_count: 0, data_and_next: None })]),
25            current_key: 0,
26        }
27    }
28
29    /// Retains the current epoch and returns the key for it.  The caller must later call `release`.
30    pub fn retain(&mut self) -> usize {
31        self.bin[self.current_key].ref_count += 1;
32        self.current_key
33    }
34
35    /// Releases a reference on the epoch with the specified key.
36    pub fn release(&mut self, mut key: usize) {
37        loop {
38            let entry = &mut self.bin[key];
39            entry.ref_count -= 1;
40            // If the reference count has reached zero and there's something to drop, drop it now.
41            // This entry holds a reference to the next entry, so we loop round and drop that
42            // reference.
43            match entry.data_and_next {
44                Some((_, next)) if entry.ref_count == 0 => {
45                    self.bin.remove(key);
46                    key = next;
47                }
48                _ => return,
49            }
50        }
51    }
52
53    /// Adds something to be dropped when all previous references are released.
54    pub fn add(&mut self, data: T) {
55        if self.bin[self.current_key].ref_count == 0 {
56            return;
57        }
58        let next = self.bin.insert(BinEntry {
59            ref_count: 1, // The current entry keeps a reference to the next one.
60            data_and_next: None,
61        });
62        self.bin[self.current_key].data_and_next = Some((data, next));
63        self.current_key = next;
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::Bin;
70    use std::sync::Arc;
71
72    #[test]
73    fn test_data_dropped_immediately_when_no_references() {
74        let data = Arc::new(());
75        let mut bin = Bin::new();
76        bin.add(data.clone());
77        assert_eq!(Arc::strong_count(&data), 1);
78    }
79
80    #[test]
81    fn test_data_dropped_when_reference_count_reaches_zero() {
82        let data = Arc::new(());
83        let mut bin = Bin::new();
84        let key = bin.retain();
85        bin.add(data.clone());
86        assert_eq!(Arc::strong_count(&data), 2);
87        bin.release(key);
88        assert_eq!(Arc::strong_count(&data), 1);
89    }
90
91    #[test]
92    fn test_data_in_later_epoch_dropped_after_earlier_epochs_are_dropped() {
93        let data1 = Arc::new(());
94        let data2 = Arc::new(());
95        let mut bin = Bin::new();
96        let key1 = bin.retain();
97        bin.add(data1.clone());
98        let key2 = bin.retain();
99        bin.add(data2.clone());
100        assert_eq!(Arc::strong_count(&data1), 2);
101        assert_eq!(Arc::strong_count(&data2), 2);
102        bin.release(key2);
103        // key2 comes after key1 so nothing should have been dropped.
104        assert_eq!(Arc::strong_count(&data1), 2);
105        assert_eq!(Arc::strong_count(&data2), 2);
106        bin.release(key1);
107        // Both should now get dropped.
108        assert_eq!(Arc::strong_count(&data1), 1);
109        assert_eq!(Arc::strong_count(&data2), 1);
110    }
111}