agentos_runtime/
accounting.rs

1//! Hierarchical count/byte admission with exact RAII ownership.
2
3use std::collections::BTreeMap;
4use std::fmt;
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::sync::{Arc, Mutex};
7
8use crate::metrics::{BufferMetricClass, ResourceMetricClass, RuntimeMetrics};
9
10/// Low-cardinality resource classes used by the sidecar admission policy.
11#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
12pub enum ResourceClass {
13    Capabilities,
14    ReadyHandles,
15    Sockets,
16    Connections,
17    BufferedBytes,
18    Datagrams,
19    HandleCommands,
20    HandleCommandBytes,
21    BridgeCalls,
22    BridgeRequestBytes,
23    BridgeResponseBytes,
24    AsyncCompletions,
25    AsyncCompletionBytes,
26    UdpDatagrams,
27    UdpBytes,
28    TlsBytes,
29    Timers,
30    Tasks,
31    ExecutorSlots,
32    ExecutorBytes,
33    Http2Connections,
34    Http2Streams,
35    Http2BufferedBytes,
36    Http2HeaderBytes,
37    Http2DataBytes,
38    Http2Commands,
39    Http2CommandBytes,
40    Http2Events,
41    Http2EventBytes,
42}
43
44impl ResourceClass {
45    pub const ALL: [Self; 29] = [
46        Self::Capabilities,
47        Self::ReadyHandles,
48        Self::Sockets,
49        Self::Connections,
50        Self::BufferedBytes,
51        Self::Datagrams,
52        Self::HandleCommands,
53        Self::HandleCommandBytes,
54        Self::BridgeCalls,
55        Self::BridgeRequestBytes,
56        Self::BridgeResponseBytes,
57        Self::AsyncCompletions,
58        Self::AsyncCompletionBytes,
59        Self::UdpDatagrams,
60        Self::UdpBytes,
61        Self::TlsBytes,
62        Self::Timers,
63        Self::Tasks,
64        Self::ExecutorSlots,
65        Self::ExecutorBytes,
66        Self::Http2Connections,
67        Self::Http2Streams,
68        Self::Http2BufferedBytes,
69        Self::Http2HeaderBytes,
70        Self::Http2DataBytes,
71        Self::Http2Commands,
72        Self::Http2CommandBytes,
73        Self::Http2Events,
74        Self::Http2EventBytes,
75    ];
76
77    pub const fn name(self) -> &'static str {
78        match self {
79            Self::Capabilities => "capabilities",
80            Self::ReadyHandles => "readyHandles",
81            Self::Sockets => "sockets",
82            Self::Connections => "connections",
83            Self::BufferedBytes => "bufferedBytes",
84            Self::Datagrams => "datagrams",
85            Self::HandleCommands => "handleCommands",
86            Self::HandleCommandBytes => "handleCommandBytes",
87            Self::BridgeCalls => "bridgeCalls",
88            Self::BridgeRequestBytes => "bridgeRequestBytes",
89            Self::BridgeResponseBytes => "bridgeResponseBytes",
90            Self::AsyncCompletions => "asyncCompletions",
91            Self::AsyncCompletionBytes => "asyncCompletionBytes",
92            Self::UdpDatagrams => "udpDatagrams",
93            Self::UdpBytes => "udpBytes",
94            Self::TlsBytes => "tlsBytes",
95            Self::Timers => "timers",
96            Self::Tasks => "tasks",
97            Self::ExecutorSlots => "executorSlots",
98            Self::ExecutorBytes => "executorBytes",
99            Self::Http2Connections => "http2Connections",
100            Self::Http2Streams => "http2Streams",
101            Self::Http2BufferedBytes => "http2BufferedBytes",
102            Self::Http2HeaderBytes => "http2HeaderBytes",
103            Self::Http2DataBytes => "http2DataBytes",
104            Self::Http2Commands => "http2Commands",
105            Self::Http2CommandBytes => "http2CommandBytes",
106            Self::Http2Events => "http2Events",
107            Self::Http2EventBytes => "http2EventBytes",
108        }
109    }
110}
111
112#[derive(Clone, Debug, Eq, PartialEq)]
113pub struct ResourceLimit {
114    pub maximum: usize,
115    pub config_path: String,
116}
117
118impl ResourceLimit {
119    pub fn new(maximum: usize, config_path: impl Into<String>) -> Self {
120        Self {
121            maximum,
122            config_path: config_path.into(),
123        }
124    }
125}
126
127#[derive(Clone, Debug, Eq, PartialEq)]
128pub struct LimitError {
129    pub scope: String,
130    pub resource: ResourceClass,
131    pub used: usize,
132    pub requested: usize,
133    pub limit: usize,
134    pub config_path: String,
135}
136
137impl fmt::Display for LimitError {
138    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
139        write!(
140            formatter,
141            "ERR_AGENTOS_RESOURCE_LIMIT: scope={} resource={} used={} requested={} limit={}; raise {}",
142            self.scope,
143            self.resource.name(),
144            self.used,
145            self.requested,
146            self.limit,
147            self.config_path
148        )
149    }
150}
151
152impl std::error::Error for LimitError {}
153
154#[derive(Clone, Debug, Eq, PartialEq)]
155pub struct ResourceUsage {
156    pub used: usize,
157    pub limit: Option<usize>,
158}
159
160#[derive(Debug, Default)]
161struct CounterState {
162    used: usize,
163    warning_active: bool,
164}
165
166#[derive(Debug)]
167struct LedgerState {
168    counters: BTreeMap<ResourceClass, CounterState>,
169}
170
171#[derive(Debug)]
172struct LedgerInner {
173    scope: String,
174    limits: BTreeMap<ResourceClass, ResourceLimit>,
175    state: Mutex<LedgerState>,
176    capacity_changed: tokio::sync::Notify,
177    integrity_failed: AtomicBool,
178    metrics: Option<RuntimeMetrics>,
179}
180
181/// One process or VM accounting scope. A child ledger reserves its parent first,
182/// so aggregate process policy and tenant policy cover the same allocation.
183#[derive(Clone, Debug)]
184pub struct ResourceLedger {
185    inner: Arc<LedgerInner>,
186    parent: Option<Arc<ResourceLedger>>,
187}
188
189impl ResourceLedger {
190    pub fn root(
191        scope: impl Into<String>,
192        limits: impl IntoIterator<Item = (ResourceClass, ResourceLimit)>,
193    ) -> Self {
194        Self::new(scope, limits, None, None)
195    }
196
197    pub fn root_with_metrics(
198        scope: impl Into<String>,
199        limits: impl IntoIterator<Item = (ResourceClass, ResourceLimit)>,
200        metrics: RuntimeMetrics,
201    ) -> Self {
202        Self::new(scope, limits, None, Some(metrics))
203    }
204
205    pub fn child(
206        scope: impl Into<String>,
207        limits: impl IntoIterator<Item = (ResourceClass, ResourceLimit)>,
208        parent: Arc<ResourceLedger>,
209    ) -> Self {
210        Self::new(scope, limits, Some(parent), None)
211    }
212
213    fn new(
214        scope: impl Into<String>,
215        limits: impl IntoIterator<Item = (ResourceClass, ResourceLimit)>,
216        parent: Option<Arc<ResourceLedger>>,
217        metrics: Option<RuntimeMetrics>,
218    ) -> Self {
219        Self {
220            inner: Arc::new(LedgerInner {
221                scope: scope.into(),
222                limits: limits.into_iter().collect(),
223                state: Mutex::new(LedgerState {
224                    counters: BTreeMap::new(),
225                }),
226                capacity_changed: tokio::sync::Notify::new(),
227                integrity_failed: AtomicBool::new(false),
228                metrics,
229            }),
230            parent,
231        }
232    }
233
234    pub fn scope(&self) -> &str {
235        &self.inner.scope
236    }
237
238    /// Reserve before allocation. A zero-sized reservation is valid and owns no
239    /// counters, which lets callers keep one unconditional cleanup path.
240    pub fn reserve(
241        &self,
242        resource: ResourceClass,
243        amount: usize,
244    ) -> Result<Reservation, LimitError> {
245        let mut allocations = Vec::with_capacity(if self.parent.is_some() { 2 } else { 1 });
246        if let Some(parent) = &self.parent {
247            parent.reserve_into(resource, amount, &mut allocations)?;
248        }
249        if let Err(error) = self.reserve_local(resource, amount) {
250            release_allocations(&mut allocations);
251            return Err(error);
252        }
253        if amount != 0 {
254            allocations.push(Allocation {
255                ledger: Arc::clone(&self.inner),
256                resource,
257                amount,
258            });
259        }
260        Ok(Reservation {
261            resource,
262            amount,
263            allocations,
264        })
265    }
266
267    fn reserve_into(
268        &self,
269        resource: ResourceClass,
270        amount: usize,
271        allocations: &mut Vec<Allocation>,
272    ) -> Result<(), LimitError> {
273        if let Some(parent) = &self.parent {
274            parent.reserve_into(resource, amount, allocations)?;
275        }
276        if let Err(error) = self.reserve_local(resource, amount) {
277            release_allocations(allocations);
278            return Err(error);
279        }
280        if amount != 0 {
281            allocations.push(Allocation {
282                ledger: Arc::clone(&self.inner),
283                resource,
284                amount,
285            });
286        }
287        Ok(())
288    }
289
290    fn reserve_local(&self, resource: ResourceClass, amount: usize) -> Result<(), LimitError> {
291        if amount == 0 {
292            return Ok(());
293        }
294        let mut state = self.inner.state.lock().unwrap_or_else(|poisoned| {
295            eprintln!(
296                "ERR_AGENTOS_RESOURCE_LEDGER_POISONED: recovering scope={} resource={}",
297                self.inner.scope,
298                resource.name()
299            );
300            poisoned.into_inner()
301        });
302        let counter = state.counters.entry(resource).or_default();
303        let requested_total = counter.used.checked_add(amount);
304        if let Some(limit) = self.inner.limits.get(&resource) {
305            if requested_total.is_none_or(|total| total > limit.maximum) {
306                return Err(LimitError {
307                    scope: self.inner.scope.clone(),
308                    resource,
309                    used: counter.used,
310                    requested: amount,
311                    limit: limit.maximum,
312                    config_path: limit.config_path.clone(),
313                });
314            }
315        }
316        counter.used = requested_total.unwrap_or(usize::MAX);
317        maybe_warn(&self.inner, resource, counter);
318        observe_usage(&self.inner, resource, counter.used);
319        Ok(())
320    }
321
322    pub fn usage(&self, resource: ResourceClass) -> ResourceUsage {
323        let state = self.inner.state.lock().unwrap_or_else(|poisoned| {
324            eprintln!(
325                "ERR_AGENTOS_RESOURCE_LEDGER_POISONED: recovering scope={} resource={}",
326                self.inner.scope,
327                resource.name()
328            );
329            poisoned.into_inner()
330        });
331        ResourceUsage {
332            used: state
333                .counters
334                .get(&resource)
335                .map_or(0, |counter| counter.used),
336            limit: self.inner.limits.get(&resource).map(|limit| limit.maximum),
337        }
338    }
339
340    pub fn configured_limit(&self, resource: ResourceClass) -> Option<ResourceLimit> {
341        self.inner.limits.get(&resource).cloned()
342    }
343
344    /// Best-effort readiness probe that does not acquire and immediately drop
345    /// a reservation. Admission must still call [`Self::reserve`]; this method
346    /// exists for POLLOUT-style hints where a transient reservation would emit
347    /// a false capacity-change notification and churn blocked producers.
348    pub fn capacity_available(&self, resource: ResourceClass, amount: usize) -> bool {
349        if let Some(parent) = &self.parent {
350            if !parent.capacity_available(resource, amount) {
351                return false;
352            }
353        }
354        if amount == 0 {
355            return true;
356        }
357        let state = self.inner.state.lock().unwrap_or_else(|poisoned| {
358            eprintln!(
359                "ERR_AGENTOS_RESOURCE_LEDGER_POISONED: recovering capacity probe scope={} resource={}",
360                self.inner.scope,
361                resource.name()
362            );
363            poisoned.into_inner()
364        });
365        let used = state
366            .counters
367            .get(&resource)
368            .map_or(0, |counter| counter.used);
369        self.inner.limits.get(&resource).is_none_or(|limit| {
370            used.checked_add(amount)
371                .is_some_and(|total| total <= limit.maximum)
372        })
373    }
374
375    pub fn is_zero(&self) -> bool {
376        let state = self.inner.state.lock().unwrap_or_else(|poisoned| {
377            eprintln!(
378                "ERR_AGENTOS_RESOURCE_LEDGER_POISONED: recovering scope={}",
379                self.inner.scope
380            );
381            poisoned.into_inner()
382        });
383        state.counters.values().all(|counter| counter.used == 0)
384    }
385
386    pub fn integrity_ok(&self) -> bool {
387        !self.inner.integrity_failed.load(Ordering::Acquire)
388    }
389
390    /// Wait without polling until any owner releases capacity. Callers must
391    /// retry their full multi-resource admission after this notification.
392    pub async fn capacity_changed(&self) {
393        let local_changed = self.inner.capacity_changed.notified();
394        if let Some(parent) = &self.parent {
395            let parent_changed = parent.inner.capacity_changed.notified();
396            tokio::select! {
397                _ = local_changed => {}
398                _ = parent_changed => {}
399            }
400        } else {
401            local_changed.await;
402        }
403    }
404
405    /// Pause an async source until the requested capacity is available. The
406    /// notification is armed before each retry so a concurrent release cannot
407    /// be missed between observing the limit and awaiting capacity.
408    pub async fn reserve_when_available(
409        &self,
410        resource: ResourceClass,
411        amount: usize,
412    ) -> Result<Reservation, LimitError> {
413        loop {
414            // Arm both accounting scopes before admission so a concurrent
415            // release cannot be missed. Ledgers are currently process -> VM;
416            // flatten this wait set if another hierarchy level is introduced.
417            let local_changed = self.inner.capacity_changed.notified();
418            let parent_changed = self
419                .parent
420                .as_ref()
421                .map(|parent| parent.inner.capacity_changed.notified());
422            match self.reserve(resource, amount) {
423                Ok(reservation) => return Ok(reservation),
424                Err(error) if amount > error.limit => return Err(error),
425                Err(_) => {
426                    if let Some(parent_changed) = parent_changed {
427                        tokio::select! {
428                            _ = local_changed => {}
429                            _ = parent_changed => {}
430                        }
431                    } else {
432                        local_changed.await;
433                    }
434                }
435            }
436        }
437    }
438}
439
440fn maybe_warn(inner: &LedgerInner, resource: ResourceClass, counter: &mut CounterState) {
441    let Some(limit) = inner.limits.get(&resource) else {
442        return;
443    };
444    // Integer comparisons avoid rounding and make zero impossible to treat as
445    // near-limit. The warning rearms only after usage falls below 70%.
446    let near =
447        counter.used != 0 && counter.used.saturating_mul(100) >= limit.maximum.saturating_mul(80);
448    if near && !counter.warning_active {
449        counter.warning_active = true;
450        eprintln!(
451            "WARN_AGENTOS_RESOURCE_NEAR_LIMIT: scope={} resource={} used={} limit={} config={}",
452            inner.scope,
453            resource.name(),
454            counter.used,
455            limit.maximum,
456            limit.config_path
457        );
458    }
459}
460
461#[derive(Debug)]
462struct Allocation {
463    ledger: Arc<LedgerInner>,
464    resource: ResourceClass,
465    amount: usize,
466}
467
468fn release_allocation(allocation: &Allocation) {
469    let mut state = allocation.ledger.state.lock().unwrap_or_else(|poisoned| {
470        eprintln!(
471            "ERR_AGENTOS_RESOURCE_LEDGER_POISONED: recovering release scope={} resource={}",
472            allocation.ledger.scope,
473            allocation.resource.name()
474        );
475        poisoned.into_inner()
476    });
477    let counter = state.counters.entry(allocation.resource).or_default();
478    if allocation.amount > counter.used {
479        eprintln!(
480            "ERR_AGENTOS_RESOURCE_ACCOUNTING_UNDERFLOW: scope={} resource={} used={} release={}",
481            allocation.ledger.scope,
482            allocation.resource.name(),
483            counter.used,
484            allocation.amount
485        );
486        allocation
487            .ledger
488            .integrity_failed
489            .store(true, Ordering::Release);
490        counter.used = 0;
491    } else {
492        counter.used -= allocation.amount;
493    }
494    if let Some(limit) = allocation.ledger.limits.get(&allocation.resource) {
495        if counter.used.saturating_mul(100) < limit.maximum.saturating_mul(70) {
496            counter.warning_active = false;
497        }
498    }
499    observe_usage(&allocation.ledger, allocation.resource, counter.used);
500    // `notify_one` retains a permit when no waiter is currently polled, which
501    // closes the release-between-retry-and-await race.
502    allocation.ledger.capacity_changed.notify_one();
503}
504
505fn observe_usage(inner: &LedgerInner, resource: ResourceClass, used: usize) {
506    let Some(metrics) = &inner.metrics else {
507        return;
508    };
509    match resource {
510        ResourceClass::Capabilities => {
511            metrics.observe_resource(ResourceMetricClass::Capabilities, used)
512        }
513        ResourceClass::ReadyHandles => {
514            metrics.observe_resource(ResourceMetricClass::ReadyHandles, used)
515        }
516        ResourceClass::Sockets => metrics.observe_resource(ResourceMetricClass::Sockets, used),
517        ResourceClass::Connections => {
518            metrics.observe_resource(ResourceMetricClass::Connections, used)
519        }
520        ResourceClass::BufferedBytes => metrics.observe_buffer(BufferMetricClass::Native, used),
521        ResourceClass::Datagrams => metrics.observe_resource(ResourceMetricClass::Datagrams, used),
522        ResourceClass::HandleCommands => {
523            metrics.observe_resource(ResourceMetricClass::HandleCommands, used)
524        }
525        ResourceClass::HandleCommandBytes => {
526            metrics.observe_buffer(BufferMetricClass::Native, used)
527        }
528        ResourceClass::BridgeCalls => {
529            metrics.observe_resource(ResourceMetricClass::BridgeCalls, used)
530        }
531        ResourceClass::BridgeRequestBytes | ResourceClass::BridgeResponseBytes => {
532            metrics.observe_buffer(BufferMetricClass::Bridge, used)
533        }
534        ResourceClass::AsyncCompletions => {
535            metrics.observe_resource(ResourceMetricClass::AsyncCompletions, used)
536        }
537        ResourceClass::AsyncCompletionBytes => {
538            metrics.observe_buffer(BufferMetricClass::Bridge, used)
539        }
540        ResourceClass::UdpDatagrams => {
541            metrics.observe_resource(ResourceMetricClass::Datagrams, used)
542        }
543        ResourceClass::UdpBytes => metrics.observe_buffer(BufferMetricClass::Datagram, used),
544        ResourceClass::TlsBytes => metrics.observe_buffer(BufferMetricClass::Tls, used),
545        ResourceClass::Timers => metrics.observe_resource(ResourceMetricClass::Timers, used),
546        ResourceClass::Tasks => metrics.observe_resource(ResourceMetricClass::Tasks, used),
547        ResourceClass::ExecutorSlots => {}
548        ResourceClass::ExecutorBytes => metrics.observe_buffer(BufferMetricClass::Executor, used),
549        ResourceClass::Http2BufferedBytes => metrics.observe_buffer(BufferMetricClass::Http2, used),
550        ResourceClass::Http2Connections => {
551            metrics.observe_resource(ResourceMetricClass::Http2Connections, used)
552        }
553        ResourceClass::Http2Streams => {
554            metrics.observe_resource(ResourceMetricClass::Http2Streams, used)
555        }
556        ResourceClass::Http2HeaderBytes
557        | ResourceClass::Http2DataBytes
558        | ResourceClass::Http2Commands
559        | ResourceClass::Http2CommandBytes
560        | ResourceClass::Http2Events
561        | ResourceClass::Http2EventBytes => {}
562    }
563}
564
565fn release_allocations(allocations: &mut Vec<Allocation>) {
566    for allocation in allocations.drain(..).rev() {
567        release_allocation(&allocation);
568    }
569}
570
571/// Exact ownership of one admitted amount. Moving the value transfers ownership;
572/// dropping it releases every scope exactly once.
573#[derive(Debug)]
574pub struct Reservation {
575    resource: ResourceClass,
576    amount: usize,
577    allocations: Vec<Allocation>,
578}
579
580/// Cloneable ownership used when a charged payload crosses an in-process
581/// response router. The underlying reservation releases only after the final
582/// envelope/consumer drops it.
583#[derive(Clone)]
584pub struct SharedReservation(Arc<Reservation>);
585
586impl SharedReservation {
587    pub fn new(reservation: Reservation) -> Self {
588        Self(Arc::new(reservation))
589    }
590
591    pub fn resource(&self) -> ResourceClass {
592        self.0.resource()
593    }
594
595    pub fn amount(&self) -> usize {
596        self.0.amount()
597    }
598}
599
600impl fmt::Debug for SharedReservation {
601    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
602        formatter
603            .debug_struct("SharedReservation")
604            .field("resource", &self.resource())
605            .field("amount", &self.amount())
606            .finish_non_exhaustive()
607    }
608}
609
610impl PartialEq for SharedReservation {
611    fn eq(&self, other: &Self) -> bool {
612        Arc::ptr_eq(&self.0, &other.0)
613    }
614}
615
616impl Eq for SharedReservation {}
617
618impl Reservation {
619    pub fn resource(&self) -> ResourceClass {
620        self.resource
621    }
622
623    pub fn amount(&self) -> usize {
624        self.amount
625    }
626
627    /// Split ownership without changing any ledger counter.
628    pub fn split(&mut self, amount: usize) -> Option<Self> {
629        if amount > self.amount {
630            return None;
631        }
632        self.amount -= amount;
633        let mut allocations = Vec::with_capacity(self.allocations.len());
634        for allocation in &mut self.allocations {
635            allocation.amount -= amount;
636            allocations.push(Allocation {
637                ledger: Arc::clone(&allocation.ledger),
638                resource: allocation.resource,
639                amount,
640            });
641        }
642        Some(Self {
643            resource: self.resource,
644            amount,
645            allocations,
646        })
647    }
648
649    /// Merge two reservations only when they refer to the same scopes and class.
650    pub fn merge(&mut self, mut other: Self) -> Result<(), Self> {
651        if self.resource != other.resource || self.allocations.len() != other.allocations.len() {
652            return Err(other);
653        }
654        if self
655            .allocations
656            .iter()
657            .zip(&other.allocations)
658            .any(|(left, right)| !Arc::ptr_eq(&left.ledger, &right.ledger))
659        {
660            return Err(other);
661        }
662        let Some(total) = self.amount.checked_add(other.amount) else {
663            return Err(other);
664        };
665        for (left, right) in self.allocations.iter_mut().zip(&other.allocations) {
666            left.amount += right.amount;
667        }
668        self.amount = total;
669        other.amount = 0;
670        other.allocations.clear();
671        Ok(())
672    }
673}
674
675impl Drop for Reservation {
676    fn drop(&mut self) {
677        release_allocations(&mut self.allocations);
678        self.amount = 0;
679    }
680}
681
682#[cfg(test)]
683mod tests {
684    use super::*;
685
686    fn limit(maximum: usize) -> [(ResourceClass, ResourceLimit); 1] {
687        [(
688            ResourceClass::BufferedBytes,
689            ResourceLimit::new(maximum, "runtime.resources.maxSocketBufferedBytes"),
690        )]
691    }
692
693    #[test]
694    fn child_reservation_charges_and_releases_both_scopes() {
695        let process = Arc::new(ResourceLedger::root("process", limit(10)));
696        let vm = ResourceLedger::child("vm-1", limit(6), Arc::clone(&process));
697        let reservation = vm.reserve(ResourceClass::BufferedBytes, 6).unwrap();
698        assert_eq!(process.usage(ResourceClass::BufferedBytes).used, 6);
699        assert_eq!(vm.usage(ResourceClass::BufferedBytes).used, 6);
700        let error = vm.reserve(ResourceClass::BufferedBytes, 1).unwrap_err();
701        assert_eq!(error.scope, "vm-1");
702        assert_eq!(process.usage(ResourceClass::BufferedBytes).used, 6);
703        drop(reservation);
704        assert!(process.is_zero());
705        assert!(vm.is_zero());
706    }
707
708    #[test]
709    fn failed_child_admission_rolls_back_parent() {
710        let process = Arc::new(ResourceLedger::root("process", limit(10)));
711        let vm = ResourceLedger::child("vm-1", limit(2), Arc::clone(&process));
712        let error = vm.reserve(ResourceClass::BufferedBytes, 3).unwrap_err();
713        assert_eq!(error.scope, "vm-1");
714        assert!(process.is_zero());
715        assert!(vm.is_zero());
716    }
717
718    #[test]
719    fn capacity_probe_checks_parent_and_child_without_changing_usage() {
720        let process = Arc::new(ResourceLedger::root("process", limit(3)));
721        let vm = ResourceLedger::child("vm-1", limit(5), Arc::clone(&process));
722        let held = process
723            .reserve(ResourceClass::BufferedBytes, 2)
724            .expect("reserve process capacity");
725
726        assert!(vm.capacity_available(ResourceClass::BufferedBytes, 1));
727        assert!(!vm.capacity_available(ResourceClass::BufferedBytes, 2));
728        assert_eq!(process.usage(ResourceClass::BufferedBytes).used, 2);
729        assert_eq!(vm.usage(ResourceClass::BufferedBytes).used, 0);
730
731        drop(held);
732        assert!(vm.capacity_available(ResourceClass::BufferedBytes, 2));
733    }
734
735    #[test]
736    fn split_and_merge_transfer_without_counter_drift() {
737        let ledger = ResourceLedger::root("vm-1", limit(10));
738        let mut source = ledger.reserve(ResourceClass::BufferedBytes, 8).unwrap();
739        let transferred = source.split(3).unwrap();
740        assert_eq!(source.amount(), 5);
741        assert_eq!(transferred.amount(), 3);
742        assert_eq!(ledger.usage(ResourceClass::BufferedBytes).used, 8);
743        source.merge(transferred).unwrap();
744        assert_eq!(source.amount(), 8);
745        assert_eq!(ledger.usage(ResourceClass::BufferedBytes).used, 8);
746        drop(source);
747        assert!(ledger.is_zero());
748    }
749
750    #[tokio::test]
751    async fn impossible_async_reservation_returns_typed_error() {
752        let ledger = ResourceLedger::root("vm-1", limit(4));
753        let error = ledger
754            .reserve_when_available(ResourceClass::BufferedBytes, 5)
755            .await
756            .unwrap_err();
757        assert_eq!(error.resource, ResourceClass::BufferedBytes);
758        assert_eq!(error.requested, 5);
759        assert_eq!(error.limit, 4);
760        assert_eq!(
761            error.config_path,
762            "runtime.resources.maxSocketBufferedBytes"
763        );
764    }
765
766    #[tokio::test]
767    async fn child_waiter_wakes_when_only_parent_capacity_changes() {
768        let process = Arc::new(ResourceLedger::root("process", limit(1)));
769        let vm = Arc::new(ResourceLedger::child(
770            "vm-1",
771            limit(2),
772            Arc::clone(&process),
773        ));
774        let held = process
775            .reserve(ResourceClass::BufferedBytes, 1)
776            .expect("fill parent");
777        let waiting_vm = Arc::clone(&vm);
778        let waiter = tokio::spawn(async move {
779            waiting_vm
780                .reserve_when_available(ResourceClass::BufferedBytes, 1)
781                .await
782        });
783        tokio::task::yield_now().await;
784        assert!(!waiter.is_finished());
785        drop(held);
786        let reservation = tokio::time::timeout(std::time::Duration::from_secs(1), waiter)
787            .await
788            .expect("parent release must wake child waiter")
789            .expect("waiter task")
790            .expect("reservation");
791        drop(reservation);
792        assert!(process.is_zero());
793        assert!(vm.is_zero());
794    }
795
796    #[test]
797    fn accounting_underflow_latches_integrity_failure() {
798        let ledger = ResourceLedger::root("vm-1", limit(1));
799        let inner = Arc::clone(&ledger.inner);
800        let malformed = Reservation {
801            resource: ResourceClass::BufferedBytes,
802            amount: 1,
803            allocations: vec![Allocation {
804                ledger: inner,
805                resource: ResourceClass::BufferedBytes,
806                amount: 1,
807            }],
808        };
809        drop(malformed);
810        assert!(!ledger.integrity_ok());
811        assert!(ledger.is_zero());
812    }
813}
agentos_runtime/accounting.rs

agentos_runtime/
accounting.rs
