mm1_address/
pool.rs

1use std::ops::Deref;
2use std::sync::{Arc, Weak};
3
4use parking_lot::Mutex;
5
6use crate::address::Address;
7use crate::subnet::{NetAddress, NetMask};
8
9mod trie;
10
11#[derive(Debug, thiserror::Error)]
12pub enum LeaseError {
13    #[error("no available addresses in the pool")]
14    Unavailable,
15}
16
17#[derive(Debug, Clone)]
18pub struct Pool {
19    shared: Arc<Mutex<Shared>>,
20}
21
22#[derive(Debug)]
23pub struct Lease {
24    net_address: NetAddress,
25    shared:      Weak<Mutex<Shared>>,
26}
27
28#[derive(Debug)]
29struct Shared {
30    main: trie::Pool,
31    used: trie::Pool,
32}
33
34impl Pool {
35    pub fn new(net_address: NetAddress) -> Self {
36        let shared = Arc::new(Mutex::new(Shared {
37            main: trie::Pool::new(net_address.address.into_u64(), net_address.mask.into_u64()),
38            used: trie::Pool::empty(),
39        }));
40        Self { shared }
41    }
42
43    pub fn lease(&self, mask: NetMask) -> Result<Lease, LeaseError> {
44        let addr = {
45            let mut shared = self.shared.lock();
46            let Shared { main, used } = &mut *shared;
47
48            if let Some(a) = main.acquire(mask.into_u64()) {
49                a
50            } else if let Some(a) = used.acquire(mask.into_u64()) {
51                std::mem::swap(main, used);
52                a
53            } else {
54                return Err(LeaseError::Unavailable);
55            }
56        };
57
58        let net_address = (Address::from_u64(addr), mask).into();
59        let shared = Arc::downgrade(&self.shared);
60        let lease = Lease {
61            net_address,
62            shared,
63        };
64        Ok(lease)
65    }
66}
67
68impl Lease {
69    pub fn net_address(&self) -> NetAddress {
70        self.net_address
71    }
72}
73
74impl Deref for Lease {
75    type Target = NetAddress;
76
77    fn deref(&self) -> &Self::Target {
78        &self.net_address
79    }
80}
81
82impl Drop for Lease {
83    fn drop(&mut self) {
84        if let Some(shared) = self.shared.upgrade() {
85            let addr = self.net_address.address.into_u64();
86            let mask = self.net_address.mask.into_u64();
87            shared.lock().used.release(addr, mask);
88        }
89    }
90}