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                // Neither trie can satisfy the request on its own, but their
55                // combined free space might. Fold `used` into `main` (which
56                // coalesces adjacent blocks) and try once more.
57                used.drain_into(main);
58                main.acquire(mask.into_u64())
59                    .ok_or(LeaseError::Unavailable)?
60            }
61        };
62
63        let net_address = (Address::from_u64(addr), mask).into();
64        let shared = Arc::downgrade(&self.shared);
65        let lease = Lease {
66            net_address,
67            shared,
68        };
69        Ok(lease)
70    }
71}
72
73impl Lease {
74    pub fn trusted(net_address: NetAddress) -> Self {
75        Self {
76            net_address,
77            shared: Weak::new(),
78        }
79    }
80
81    pub fn net_address(&self) -> NetAddress {
82        self.net_address
83    }
84}
85
86impl Deref for Lease {
87    type Target = NetAddress;
88
89    fn deref(&self) -> &Self::Target {
90        &self.net_address
91    }
92}
93
94impl Drop for Lease {
95    fn drop(&mut self) {
96        if let Some(shared) = self.shared.upgrade() {
97            let addr = self.net_address.address.into_u64();
98            let mask = self.net_address.mask.into_u64();
99            shared.lock().used.release(addr, mask);
100        }
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    fn mask(bits: u8) -> NetMask {
109        NetMask::try_from(bits).expect("valid mask")
110    }
111
112    // Regression test for #147: freeing one half of a block must not strand the
113    // other half in a different trie. A lease that needs the whole block must
114    // still succeed after the two halves are freed to different tries.
115    #[test]
116    fn lease_reuses_coalesced_halves() {
117        // A /63 pool holds exactly two /64 addresses.
118        let pool = Pool::new(NetAddress {
119            address: Address::from_u64(0),
120            mask:    mask(63),
121        });
122
123        // Take one /64 (leaves the sibling in `main`), then free it (to `used`).
124        let leased = pool.lease(mask(64)).expect("first /64");
125        drop(leased);
126
127        // The two /64 halves are now split across `main` and `used`; a /63 needs
128        // both, so it must coalesce them.
129        assert!(
130            pool.lease(mask(63)).is_ok(),
131            "a /63 lease should reuse the two freed /64 halves"
132        );
133    }
134}
mm1_address/pool.rs

mm1_address/
pool.rs
