o3 0.6.0

shared-nothing primitives
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

use std::cell::Cell;
use std::marker::{PhantomData, PhantomPinned};
use std::pin::Pin;
use std::ptr::NonNull;

use crate::marker::ThreadBound;

pub struct ByteBudget {
    limit: usize,
    used: Cell<usize>,
    _pin: PhantomPinned,
    _thread: ThreadBound,
}

impl ByteBudget {
    pub fn new(limit: usize) -> Self {
        Self {
            limit,
            used: Cell::new(0),
            _pin: PhantomPinned,
            _thread: ThreadBound::NEW,
        }
    }

    pub fn limit(&self) -> usize {
        self.limit
    }

    pub fn used(&self) -> usize {
        self.used.get()
    }

    pub fn handle<'d>(self: Pin<&'d Self>) -> ByteBudgetHandle<'d> {
        ByteBudgetHandle(NonNull::from(self.get_ref()), PhantomData)
    }

    pub fn try_acquire(self: Pin<&Self>, amount: usize) -> Option<ByteLease<'_>> {
        self.handle().try_acquire(amount)
    }
}

#[derive(Clone, Copy)]
pub struct ByteBudgetHandle<'d>(NonNull<ByteBudget>, PhantomData<&'d ByteBudget>);

impl<'d> ByteBudgetHandle<'d> {
    fn budget(self) -> &'d ByteBudget {
        unsafe { self.0.as_ref() }
    }

    pub fn limit(self) -> usize {
        self.budget().limit
    }

    pub fn used(self) -> usize {
        self.budget().used.get()
    }

    pub fn try_acquire(self, amount: usize) -> Option<ByteLease<'d>> {
        let budget = self.budget();
        let used = budget.used.get().checked_add(amount)?;
        if used > budget.limit {
            return None;
        }
        budget.used.set(used);
        Some(ByteLease {
            budget: self,
            amount,
        })
    }

    fn release(self, amount: usize) {
        let budget = self.budget();
        assert!(budget.used.get() >= amount, "byte budget underflow");
        budget.used.set(budget.used.get() - amount);
    }
}

pub struct ByteLease<'d> {
    budget: ByteBudgetHandle<'d>,
    amount: usize,
}

impl ByteLease<'_> {
    pub fn amount(&self) -> usize {
        self.amount
    }

    pub fn try_grow(&mut self, additional: usize) -> bool {
        let Some(lease) = self.budget.try_acquire(additional) else {
            return false;
        };
        self.amount += lease.amount;
        std::mem::forget(lease);
        true
    }

    pub fn shrink(&mut self, amount: usize) {
        assert!(amount <= self.amount, "byte lease underflow");
        self.amount -= amount;
        self.budget.release(amount);
    }
}

impl Drop for ByteLease<'_> {
    fn drop(&mut self) {
        self.budget.release(self.amount);
    }
}