flute/sender/
toiallocator.rs

1use std::sync::{Arc, Mutex};
2
3use rand::Rng;
4
5use crate::common::lct;
6
7use super::TOIMaxLength;
8
9#[derive(Debug)]
10struct ToiAllocatorInternal {
11    toi_reserved: std::collections::HashSet<u128>,
12    toi: u128,
13    toi_max_length: TOIMaxLength,
14}
15
16#[derive(Debug)]
17pub struct ToiAllocator {
18    internal: Mutex<ToiAllocatorInternal>,
19}
20
21/// Struct containing a TOI
22#[derive(Debug)]
23pub struct Toi {
24    allocator: Arc<ToiAllocator>,
25    value: u128,
26}
27
28impl Drop for Toi {
29    fn drop(&mut self) {
30        self.allocator.release(self.value);
31    }
32}
33
34impl Toi {
35    /// Get Value of TOI
36    pub fn get(&self) -> u128 {
37        self.value
38    }
39}
40
41impl ToiAllocatorInternal {
42    fn new(toi_max_length: TOIMaxLength, toi_initial_value: Option<u128>) -> Self {
43        let mut toi = match toi_initial_value {
44            Some(0) => 1,
45            Some(n) => n,
46            None => {
47                let mut rng = rand::rng();
48                rng.random()
49            }
50        };
51
52        toi = Self::to_max_length(toi, toi_max_length);
53        if toi == lct::TOI_FDT {
54            toi += 1;
55        }
56
57        Self {
58            toi_reserved: std::collections::HashSet::new(),
59            toi,
60            toi_max_length,
61        }
62    }
63
64    fn to_max_length(toi: u128, toi_max_length: TOIMaxLength) -> u128 {
65        match toi_max_length {
66            TOIMaxLength::ToiMax16 => toi & 0xFFFFu128,
67            TOIMaxLength::ToiMax32 => toi & 0xFFFFFFFFu128,
68            TOIMaxLength::ToiMax48 => toi & 0xFFFFFFFFFFFFu128,
69            TOIMaxLength::ToiMax64 => toi & 0xFFFFFFFFFFFFFFFFu128,
70            TOIMaxLength::ToiMax80 => toi & 0xFFFFFFFFFFFFFFFFFFFFu128,
71            TOIMaxLength::ToiMax112 => toi,
72        }
73    }
74
75    fn allocate(&mut self) -> u128 {
76        let ret = self.toi;
77        assert!(!self.toi_reserved.contains(&ret));
78        self.toi_reserved.insert(ret);
79
80        loop {
81            self.toi = Self::to_max_length(self.toi + 1, self.toi_max_length);
82            if self.toi == lct::TOI_FDT {
83                self.toi = 1;
84            }
85
86            if !self.toi_reserved.contains(&self.toi) {
87                break;
88            }
89
90            log::warn!("TOI {} is already used by a file or reserved", self.toi)
91        }
92        ret
93    }
94
95    fn release(&mut self, toi: u128) {
96        let success = self.toi_reserved.remove(&toi);
97        debug_assert!(success);
98    }
99}
100
101impl ToiAllocator {
102    pub fn new(toi_max_length: TOIMaxLength, toi_initial_value: Option<u128>) -> Arc<Self> {
103        Arc::new(Self {
104            internal: Mutex::new(ToiAllocatorInternal::new(toi_max_length, toi_initial_value)),
105        })
106    }
107
108    pub fn allocate(allocator: &Arc<Self>) -> Box<Toi> {
109        let mut db = allocator.internal.lock().unwrap();
110        let toi = db.allocate();
111        Box::new(Toi {
112            allocator: allocator.clone(),
113            value: toi,
114        })
115    }
116
117    pub fn allocate_toi_fdt(allocator: &Arc<Self>) -> Box<Toi> {
118        Box::new(Toi {
119            allocator: allocator.clone(),
120            value: 0,
121        })
122    }
123
124    pub fn release(&self, toi: u128) {
125        if toi == lct::TOI_FDT {
126            return;
127        }
128        {
129            let mut db = self.internal.lock().unwrap();
130            db.release(toi);
131        }
132    }
133}