1use std::collections::HashMap;
2use std::sync::{Mutex, MutexGuard};
3use std::time::{SystemTime, UNIX_EPOCH};
4
5use super::{
6 budget_commit_metadata, checked_committed_cost_units, validate_optional_budget_identity,
7 ApprovalRequiredBudgetHold, AuthorizedBudgetHold, BudgetAdmissionBinding,
8 BudgetAuthorizationOutcome, BudgetAuthorizeCumulativeApprovalRequest,
9 BudgetAuthorizeHoldDecision, BudgetAuthorizeHoldRequest,
10 BudgetCancelCapturedBeforeDispatchRequest, BudgetCaptureHoldDecision, BudgetCaptureHoldRequest,
11 BudgetCaptureInvocationRequest, BudgetCapturedBeforeDispatchCancellationDecision,
12 BudgetCumulativeApprovalAccountKey, BudgetCumulativeApprovalAccountUsage,
13 BudgetCumulativeApprovalAuthorizationDecision, BudgetCumulativeApprovalMutation,
14 BudgetCumulativeApprovalRequest, BudgetCumulativeApprovalState, BudgetCumulativeApprovalUsage,
15 BudgetEventAuthority, BudgetHoldDispositionView, BudgetHoldMutationDecision,
16 BudgetHoldSnapshot, BudgetInvocationCaptureDecision, BudgetInvocationQuota,
17 BudgetInvocationQuotaMutation, BudgetInvocationQuotaUsage, BudgetInvocationState,
18 BudgetMonetaryState, BudgetMutationKind, BudgetMutationRecord, BudgetQuotaKey,
19 BudgetReconcileHoldDecision, BudgetReconcileHoldRequest, BudgetReleaseHoldDecision,
20 BudgetReleaseHoldRequest, BudgetReverseHoldDecision, BudgetReverseHoldRequest, BudgetStore,
21 BudgetStoreError, BudgetUsageRecord, DeniedBudgetHold, ReservedHoldEnvelope,
22};
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25struct BudgetHoldState {
26 capability_id: String,
27 grant_index: usize,
28 admission_binding: Option<BudgetAdmissionBinding>,
29 authorized_exposure_units: u64,
30 remaining_exposure_units: u64,
31 invocation_state: BudgetInvocationState,
32 invocation_quotas: Vec<BudgetInvocationQuota>,
33 legacy_captured_invocation_quota: Option<BudgetInvocationQuota>,
34 captured_cancellation_allowed: bool,
35 cumulative_approval: Option<BudgetCumulativeApprovalHoldState>,
36 monetary_state: BudgetMonetaryState,
37 authority: Option<BudgetEventAuthority>,
38 reserved_until: Option<i64>,
39 reserved_currency: Option<String>,
40 reserved_payment_reference: Option<String>,
41 reserved_envelope: ReservedHoldEnvelope,
42}
43
44#[derive(Debug, Clone, PartialEq, Eq)]
45struct BudgetCumulativeApprovalHoldState {
46 request: BudgetCumulativeApprovalRequest,
47 state: BudgetCumulativeApprovalState,
48 approval_set_digest: Option<String>,
49}
50
51#[derive(Debug, Clone, PartialEq, Eq)]
52struct BudgetInvocationQuotaState {
53 max_invocations: u32,
54 reserved_invocations: u32,
55 captured_invocations: u32,
56}
57
58#[derive(Debug, Clone, PartialEq, Eq)]
59struct BudgetCumulativeApprovalAccountState {
60 authority_threshold_units: u64,
61 reserved_authorized_units: u64,
62 captured_authorized_units: u64,
63 version: u64,
64}
65
66#[derive(Debug, Clone, PartialEq, Eq)]
67enum BudgetMutationRequest {
68 Increment {
69 capability_id: String,
70 grant_index: usize,
71 max_invocations: Option<u32>,
72 },
73 Authorize {
74 capability_id: String,
75 grant_index: usize,
76 hold_id: Option<String>,
77 cost_units: u64,
78 max_invocations: Option<u32>,
79 max_cost_per_invocation: Option<u64>,
80 max_total_cost_units: Option<u64>,
81 authority: Option<BudgetEventAuthority>,
82 },
83 AuthorizeComposite(Box<BudgetAuthorizeHoldRequest>),
84 AuthorizeCumulativeApproval(Box<BudgetAuthorizeCumulativeApprovalRequest>),
85 CaptureInvocation {
86 capability_id: String,
87 grant_index: usize,
88 hold_id: String,
89 trusted_time: Option<u64>,
90 authority: Option<BudgetEventAuthority>,
91 },
92 CancelCapturedBeforeDispatch {
93 capability_id: String,
94 grant_index: usize,
95 hold_id: String,
96 authority: Option<BudgetEventAuthority>,
97 },
98 Reverse {
99 capability_id: String,
100 grant_index: usize,
101 hold_id: Option<String>,
102 cost_units: u64,
103 expected_cumulative_approval_state: Option<BudgetCumulativeApprovalState>,
104 authority: Option<BudgetEventAuthority>,
105 },
106 Release {
107 capability_id: String,
108 grant_index: usize,
109 hold_id: Option<String>,
110 cost_units: u64,
111 authority: Option<BudgetEventAuthority>,
112 },
113 Reconcile {
114 capability_id: String,
115 grant_index: usize,
116 hold_id: Option<String>,
117 exposed_cost_units: u64,
118 realized_cost_units: u64,
119 authority: Option<BudgetEventAuthority>,
120 },
121 CaptureMonetary(BudgetCaptureHoldRequest),
122}
123
124#[derive(Debug, Clone)]
125struct RecordedBudgetMutation {
126 request: BudgetMutationRequest,
127 record: BudgetMutationRecord,
128}
129
130pub struct InMemoryBudgetStore {
131 inner: Mutex<InMemoryBudgetStoreInner>,
132}
133
134impl Default for InMemoryBudgetStore {
135 fn default() -> Self {
136 Self {
137 inner: Mutex::new(InMemoryBudgetStoreInner::default()),
138 }
139 }
140}
141
142impl InMemoryBudgetStore {
143 pub fn new() -> Self {
144 Self::default()
145 }
146
147 fn lock_inner(&self) -> Result<MutexGuard<'_, InMemoryBudgetStoreInner>, BudgetStoreError> {
148 self.inner.lock().map_err(|_| {
149 BudgetStoreError::Invariant("in-memory budget store lock poisoned".to_string())
150 })
151 }
152
153 fn recorded_mutation_decision(
154 &self,
155 inner: &InMemoryBudgetStoreInner,
156 event_id: Option<&str>,
157 ) -> Result<BudgetHoldMutationDecision, BudgetStoreError> {
158 let event = match event_id {
159 Some(event_id) => inner.events.iter().find(|event| event.event_id == event_id),
160 None => inner.events.last(),
161 }
162 .cloned()
163 .ok_or_else(|| {
164 BudgetStoreError::Invariant(
165 "mutation event disappeared while building decision".to_string(),
166 )
167 })?;
168 Ok(BudgetHoldMutationDecision {
169 hold_id: event.hold_id,
170 admission_binding: event.admission_binding,
171 exposure_units: event.exposure_units,
172 realized_spend_units: event.realized_spend_units,
173 committed_cost_units_after: checked_committed_cost_units(
174 event.total_cost_exposed_after,
175 event.total_cost_realized_spend_after,
176 )?,
177 invocation_count_after: event.invocation_count_after,
178 invocation_quota_usages: event.invocation_quota_usages,
179 cumulative_approval: event.cumulative_approval,
180 invocation_state: event.invocation_state_after,
181 monetary_state: event.monetary_state_after,
182 metadata: budget_commit_metadata(
183 self,
184 event.authority,
185 Some(event.event_seq),
186 Some(event.event_id),
187 Some(event.recorded_at),
188 ),
189 })
190 }
191}
192
193#[derive(Default)]
194struct InMemoryBudgetStoreInner {
195 counts: HashMap<(String, usize), BudgetUsageRecord>,
196 events: Vec<BudgetMutationRecord>,
197 explicit_events: HashMap<String, RecordedBudgetMutation>,
198 holds: HashMap<String, BudgetHoldState>,
199 hold_authorizations: HashMap<String, BudgetMutationRequest>,
200 operation_authorizations: HashMap<String, BudgetMutationRequest>,
201 invocation_quotas: HashMap<BudgetQuotaKey, BudgetInvocationQuotaState>,
202 legacy_reversible_invocations: HashMap<(String, usize), u32>,
203 cumulative_approval_accounts:
204 HashMap<BudgetCumulativeApprovalAccountKey, BudgetCumulativeApprovalAccountState>,
205 next_seq: u64,
206}
207
208include!("in_memory/composite.rs");
209include!("in_memory/admission.rs");
210include!("in_memory/terminal.rs");
211include!("in_memory/trait_impl.rs");
212
213fn unix_now() -> i64 {
214 SystemTime::now()
215 .duration_since(UNIX_EPOCH)
216 .map(|duration| duration.as_secs() as i64)
217 .unwrap_or(0)
218}