Skip to main content

monoloop_loop/transaction/
tool_capacity.rs

1//! Bounded global, per-transaction, and per-tool concurrency for linked tools.
2
3use monoloop_contracts::ToolId;
4use std::collections::HashMap;
5use std::sync::atomic::{AtomicUsize, Ordering};
6use std::sync::{Arc, Mutex};
7
8/// Shared process-wide concurrent tool execution limit.
9#[derive(Debug)]
10pub struct SharedToolCapacity {
11    max: usize,
12    active: AtomicUsize,
13}
14
15impl SharedToolCapacity {
16    /// Create with a maximum concurrent executions across all transactions.
17    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    /// Unlimited (practical) capacity for isolated tests.
25    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    /// Current active count.
50    pub fn active(&self) -> usize {
51        self.active.load(Ordering::SeqCst)
52    }
53}
54
55/// Per-transaction tool capacity tracker (item queue + concurrency + per-tool).
56#[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
72/// RAII permit for one running tool execution.
73pub 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    /// Build for one transaction.
93    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    /// Register per-tool max concurrency from a resolved set.
109    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    /// Reserve a queue slot before validation work (released on failure or after acquire).
121    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    /// Drop a prior queue reservation without starting.
138    pub fn dequeue(self: &Arc<Self>) {
139        self.txn_queued.fetch_sub(1, Ordering::SeqCst);
140    }
141
142    /// Move from queued to running if capacity allows.
143    pub fn try_acquire(self: &Arc<Self>, tool_id: &ToolId) -> Option<ToolPermit> {
144        // Per-tool + txn concurrency.
145        {
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    /// Active executions in this transaction.
173    pub fn active(&self) -> usize {
174        self.txn_active.load(Ordering::SeqCst)
175    }
176
177    /// Queued starts waiting for concurrency.
178    pub fn queued(&self) -> usize {
179        self.txn_queued.load(Ordering::SeqCst)
180    }
181}