monoloop_loop/transaction/
tool_capacity.rs1use monoloop_contracts::ToolId;
4use std::collections::HashMap;
5use std::sync::atomic::{AtomicUsize, Ordering};
6use std::sync::{Arc, Mutex};
7
8#[derive(Debug)]
10pub struct SharedToolCapacity {
11 max: usize,
12 active: AtomicUsize,
13}
14
15impl SharedToolCapacity {
16 pub fn new(max: usize) -> Arc<Self> {
18 Arc::new(Self {
19 max: max.max(1),
20 active: AtomicUsize::new(0),
21 })
22 }
23
24 pub fn unlimited() -> Arc<Self> {
26 Self::new(usize::MAX / 4)
27 }
28
29 fn try_acquire(&self) -> bool {
30 loop {
31 let cur = self.active.load(Ordering::SeqCst);
32 if cur >= self.max {
33 return false;
34 }
35 if self
36 .active
37 .compare_exchange(cur, cur + 1, Ordering::SeqCst, Ordering::SeqCst)
38 .is_ok()
39 {
40 return true;
41 }
42 }
43 }
44
45 fn release(&self) {
46 self.active.fetch_sub(1, Ordering::SeqCst);
47 }
48
49 pub fn active(&self) -> usize {
51 self.active.load(Ordering::SeqCst)
52 }
53}
54
55#[derive(Debug)]
57pub struct TransactionToolCapacity {
58 shared: Arc<SharedToolCapacity>,
59 max_concurrent: usize,
60 max_queued: usize,
61 txn_active: AtomicUsize,
62 txn_queued: AtomicUsize,
63 per_tool: Mutex<HashMap<ToolId, ToolSlot>>,
64}
65
66#[derive(Debug, Default)]
67struct ToolSlot {
68 max_concurrent: usize,
69 active: usize,
70}
71
72pub struct ToolPermit {
74 shared: Arc<SharedToolCapacity>,
75 txn: Arc<TransactionToolCapacity>,
76 tool_id: ToolId,
77}
78
79impl Drop for ToolPermit {
80 fn drop(&mut self) {
81 self.shared.release();
82 self.txn.txn_active.fetch_sub(1, Ordering::SeqCst);
83 if let Ok(mut map) = self.txn.per_tool.lock() {
84 if let Some(slot) = map.get_mut(&self.tool_id) {
85 slot.active = slot.active.saturating_sub(1);
86 }
87 }
88 }
89}
90
91impl TransactionToolCapacity {
92 pub fn new(
94 shared: Arc<SharedToolCapacity>,
95 max_concurrent: usize,
96 max_queued: usize,
97 ) -> Arc<Self> {
98 Arc::new(Self {
99 shared,
100 max_concurrent: max_concurrent.max(1),
101 max_queued: max_queued.max(1),
102 txn_active: AtomicUsize::new(0),
103 txn_queued: AtomicUsize::new(0),
104 per_tool: Mutex::new(HashMap::new()),
105 })
106 }
107
108 pub fn configure_tool(self: &Arc<Self>, tool_id: ToolId, max_concurrent: usize) {
110 let mut map = self.per_tool.lock().unwrap_or_else(|e| e.into_inner());
111 map.insert(
112 tool_id,
113 ToolSlot {
114 max_concurrent: max_concurrent.max(1),
115 active: 0,
116 },
117 );
118 }
119
120 pub fn try_enqueue(self: &Arc<Self>) -> bool {
122 loop {
123 let q = self.txn_queued.load(Ordering::SeqCst);
124 if q >= self.max_queued {
125 return false;
126 }
127 if self
128 .txn_queued
129 .compare_exchange(q, q + 1, Ordering::SeqCst, Ordering::SeqCst)
130 .is_ok()
131 {
132 return true;
133 }
134 }
135 }
136
137 pub fn dequeue(self: &Arc<Self>) {
139 self.txn_queued.fetch_sub(1, Ordering::SeqCst);
140 }
141
142 pub fn try_acquire(self: &Arc<Self>, tool_id: &ToolId) -> Option<ToolPermit> {
144 {
146 let mut map = self.per_tool.lock().unwrap_or_else(|e| e.into_inner());
147 let slot = map.entry(tool_id.clone()).or_insert(ToolSlot {
148 max_concurrent: 1,
149 active: 0,
150 });
151 if slot.active >= slot.max_concurrent {
152 return None;
153 }
154 let txn = self.txn_active.load(Ordering::SeqCst);
155 if txn >= self.max_concurrent {
156 return None;
157 }
158 if !self.shared.try_acquire() {
159 return None;
160 }
161 slot.active += 1;
162 self.txn_active.fetch_add(1, Ordering::SeqCst);
163 self.txn_queued.fetch_sub(1, Ordering::SeqCst);
164 }
165 Some(ToolPermit {
166 shared: Arc::clone(&self.shared),
167 txn: Arc::clone(self),
168 tool_id: tool_id.clone(),
169 })
170 }
171
172 pub fn active(&self) -> usize {
174 self.txn_active.load(Ordering::SeqCst)
175 }
176
177 pub fn queued(&self) -> usize {
179 self.txn_queued.load(Ordering::SeqCst)
180 }
181}