1#![deny(missing_docs)]
31
32pub(crate) mod instant;
33
34#[cfg(feature = "test-support")]
35pub mod test_support;
36
37use instant::StoredInstant;
38use std::num::{NonZeroI64, NonZeroUsize};
39use std::sync::Arc;
40
41use async_trait::async_trait;
42use jiff::{SignedDuration, Timestamp};
43use serde::{Deserialize, Serialize};
44use sqlx::postgres::PgPoolOptions;
45use sqlx::{PgPool, Postgres, Row, Transaction};
46use tokio::sync::broadcast;
47
48use tollgate_core::{
49 AccountId, AccountSnapshot, AccountStatus, BudgetSchedule, BudgetView, CapacityClass,
50 CostTable, CostUnits, EnforcementMode, FencingToken, Generation, KeyId, LeaseGrant, LeaseId,
51 Period, PermissionBits, PolicyRevision, Principal, PublishableSnapshot, ResolvedLimits,
52 Rollover, UsageEvent,
53};
54use tollgate_store::{
55 AccountConfig, AccountView, AdminReceipt, AdminState, AdminStore, AllocateError, Allocation,
56 BudgetError, Conservation, CreateAccountError, GrantPolicy, IngestError, IngestReport,
57 KeyDirectory, KeyError, KeyRecord, KeySnapshotError, KeySummary, LeaseAllocator,
58 PUSH_CHANNEL_CAPACITY, PublishSnapshotError, ReclaimBatch, ReclaimedLease, Revocation,
59 RolledAccount, RolloverBatch, SetStatusError, SnapshotPush, SnapshotResolution, SnapshotSource,
60 StatusChange, StoreError, StoreHealth, UsageSink, pushes_exceed_capacity,
61 validate_key_page_limit,
62};
63
64const STATE_ACTIVE: i16 = 0;
65const STATE_RELEASED: i16 = 1;
66const STATE_EXPIRED: i16 = 2;
67
68#[derive(Debug, Clone, Copy)]
72struct StoredId(u128);
73
74impl Serialize for StoredId {
75 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
76 where
77 S: serde::Serializer,
78 {
79 match u64::try_from(self.0) {
80 Ok(value) => serializer.serialize_u64(value),
81 Err(_) => serializer.collect_str(&format_args!("{:032x}", self.0)),
82 }
83 }
84}
85
86impl<'de> Deserialize<'de> for StoredId {
87 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
88 where
89 D: serde::Deserializer<'de>,
90 {
91 struct StoredIdVisitor;
92
93 impl serde::de::Visitor<'_> for StoredIdVisitor {
94 type Value = StoredId;
95
96 fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97 formatter.write_str("a legacy u64 number or a canonical 128-bit identifier string")
98 }
99
100 fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
101 Ok(StoredId(u128::from(value)))
102 }
103
104 fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
105 where
106 E: serde::de::Error,
107 {
108 value
109 .parse::<AccountId>()
110 .map(|id| StoredId(id.0))
111 .map_err(E::custom)
112 }
113 }
114
115 deserializer.deserialize_any(StoredIdVisitor)
116 }
117}
118
119fn digest_from(bytes: &[u8], key_id: KeyId) -> Result<[u8; 32], StoreError> {
124 <[u8; 32]>::try_from(bytes).map_err(|_| {
125 StoreError(format!(
126 "credential {key_id} has a {}-byte digest in storage, expected 32",
127 bytes.len()
128 ))
129 })
130}
131
132#[async_trait]
133impl KeyDirectory for PostgresStore {
134 async fn credential_activity(
135 &self,
136 keys: &[KeyId],
137 ) -> Result<Vec<tollgate_store::CredentialActivity>, StoreError> {
138 use tollgate_store::{CredentialActivity, CredentialActivityState};
139 let mut result = Vec::with_capacity(keys.len());
140 for chunk in keys.chunks(tollgate_store::MAX_INGEST_BATCH) {
143 let ids: Vec<_> = chunk.iter().map(|id| id_bytes(id.0)).collect();
144 let rows = sqlx::query(
145 "SELECT evidence.key_id IS NOT NULL, evidence.last_committed_at_us
149 FROM UNNEST($1::bytea[]) WITH ORDINALITY AS requested(key_id, ordinal)
150 LEFT JOIN LATERAL (
151 SELECT k.key_id, a.last_committed_at_us
152 FROM tollgate_credential_keys k
153 LEFT JOIN tollgate_credential_activity a ON a.key_id = k.key_id
154 WHERE k.key_id = requested.key_id LIMIT 1
155 ) AS evidence ON true
156 ORDER BY requested.ordinal",
157 )
158 .bind(&ids)
159 .fetch_all(&self.pool)
160 .await
161 .map_err(storage)?;
162 if rows.len() != chunk.len() {
163 return Err(StoreError("incomplete credential activity read".into()));
164 }
165 for (&key_id, row) in chunk.iter().zip(rows) {
166 let state = if !row.get::<bool, _>(0) {
167 CredentialActivityState::Unknown
168 } else if let Some(at) = row.get::<Option<i64>, _>(1) {
169 CredentialActivityState::Committed {
170 last_committed_at: micros_ts(at, "credential activity")?,
171 }
172 } else {
173 CredentialActivityState::Unobserved
174 };
175 result.push(CredentialActivity { key_id, state });
176 }
177 }
178 Ok(result)
179 }
180
181 async fn insert_key(&self, record: KeyRecord) -> Result<(), KeyError> {
182 self.insert_credential(record, None)
183 .await
184 .map(|receipt| receipt.outcome)
185 }
186
187 async fn revoke_key(&self, key_id: KeyId, now: Timestamp) -> Result<Revocation, KeyError> {
188 self.revoke_key_audited(key_id, now)
189 .await
190 .map(|receipt| receipt.outcome)
191 }
192
193 async fn revoke_key_audited(
194 &self,
195 key_id: KeyId,
196 now: Timestamp,
197 ) -> Result<AdminReceipt<Revocation>, KeyError> {
198 let mut tx = self.pool.begin().await.map_err(storage)?;
199 let result = async {
200 let row = sqlx::query(
204 "SELECT account_id, revoked_at_us IS NOT NULL
205 FROM tollgate_credential_keys WHERE key_id = $1 FOR UPDATE",
206 )
207 .bind(id_bytes(key_id.0))
208 .fetch_optional(&mut *tx)
209 .await
210 .map_err(storage)?
211 .ok_or(KeyError::UnknownKey)?;
212 let account_bytes: Vec<u8> = row.get(0);
213 let account_bytes: [u8; 16] = account_bytes
214 .try_into()
215 .map_err(|_| StoreError("credential account identifier is not 16 bytes".into()))?;
216 let account_id = AccountId(u128::from_be_bytes(account_bytes));
217 let revoked: bool = row.get(1);
218 let before = AdminState::Credential {
219 account_id,
220 key_id,
221 revoked,
222 };
223 let outcome = if revoked {
224 Revocation::AlreadyRetired
225 } else {
226 sqlx::query(
227 "UPDATE tollgate_credential_keys SET revoked_at_us = $2 WHERE key_id = $1",
228 )
229 .bind(id_bytes(key_id.0))
230 .bind(ts_micros(now))
231 .execute(&mut *tx)
232 .await
233 .map_err(storage)?;
234 Revocation::Retired
235 };
236 Ok(AdminReceipt::new(
237 outcome,
238 before,
239 AdminState::Credential {
240 account_id,
241 key_id,
242 revoked: true,
243 },
244 ))
245 }
246 .await;
247 finish_transaction(tx, result).await
248 }
249
250 async fn publish_key_snapshot(
251 &self,
252 account: AccountId,
253 key: KeyId,
254 snapshot: PublishableSnapshot,
255 ) -> Result<AdminReceipt<()>, KeySnapshotError> {
256 let generation = i64::try_from(snapshot.generation.0).map_err(|_| {
257 StoreError("snapshot generation exceeds PostgreSQL BIGINT range".into())
258 })?;
259 let mut tx = self.pool.begin().await.map_err(storage)?;
260 let result = async {
261 let (principal, revoked) = lock_account_key(&mut tx, account, key).await?;
262 if revoked {
263 return Err(KeySnapshotError::Retired { key_id: key });
264 }
265 if snapshot.key_id != Some(key) {
266 return Err(PublishSnapshotError::CredentialMismatch { key_id: key }.into());
267 }
268 let published = publish_in_tx(&mut tx, principal, generation, snapshot).await?;
269 Ok((principal, published))
270 }
271 .await;
272 let (principal, (written, published, before, after)) =
273 finish_transaction(tx, result).await?;
274 if written {
275 self.push_to_subscribers(SnapshotPush {
276 principal,
277 resolution: SnapshotResolution::Present(published),
278 });
279 }
280 Ok(AdminReceipt::new((), before, after))
281 }
282
283 async fn remove_key_snapshot(
284 &self,
285 account: AccountId,
286 key: KeyId,
287 ) -> Result<AdminReceipt<()>, KeySnapshotError> {
288 let mut tx = self.pool.begin().await.map_err(storage)?;
289 let result = async {
290 let (principal, _revoked) = lock_account_key(&mut tx, account, key).await?;
291 Ok::<_, KeySnapshotError>((principal, remove_in_tx(&mut tx, principal).await?))
292 }
293 .await;
294 let (principal, receipt) = finish_transaction(tx, result).await?;
295 self.announce_removal(principal, &receipt);
296 Ok(receipt)
297 }
298
299 async fn active_keys(&self, now: Timestamp) -> Result<Vec<KeyRecord>, StoreError> {
300 let cutoff = StoredInstant::from(now);
301 let rows = sqlx::query(
307 "SELECT key_id, account_id, principal, digest, not_after_floor_us, not_after_submicro_ns, not_after_is_lower_bound
308 FROM tollgate_credential_keys
309 WHERE revoked_at_us IS NULL
310 AND (not_after_floor_us IS NULL OR not_after_submicro_ns IS NULL
311 OR (not_after_floor_us, not_after_submicro_ns) > ($1, $2))
312 ORDER BY key_id",
313 )
314 .bind(cutoff.micros)
315 .bind(cutoff.submicro_nanos)
316 .fetch_all(&self.pool)
317 .await
318 .map_err(storage)?;
319
320 rows.into_iter().map(credential_from_row).collect()
321 }
322
323 async fn account_keys(
324 &self,
325 account: AccountId,
326 after: Option<KeyId>,
327 limit: NonZeroUsize,
328 ) -> Result<Vec<KeySummary>, StoreError> {
329 validate_key_page_limit(limit)?;
330 let sql = if after.is_some() {
333 "SELECT key_id, not_after_floor_us, not_after_submicro_ns,
334 not_after_is_lower_bound, revoked_at_us
335 FROM tollgate_credential_keys
336 WHERE account_id = $1 AND key_id > $3
337 ORDER BY key_id LIMIT $2"
338 } else {
339 "SELECT key_id, not_after_floor_us, not_after_submicro_ns,
340 not_after_is_lower_bound, revoked_at_us
341 FROM tollgate_credential_keys
342 WHERE account_id = $1
343 ORDER BY key_id LIMIT $2"
344 };
345 let mut query = sqlx::query(sql)
346 .bind(id_bytes(account.0))
347 .bind(i64::try_from(limit.get()).unwrap_or(i64::MAX));
348 if let Some(cursor) = after {
349 query = query.bind(id_bytes(cursor.0));
350 }
351 let rows = query.fetch_all(&self.pool).await.map_err(storage)?;
352 rows.into_iter().map(summary_from_row).collect()
353 }
354
355 async fn insert_key_within(
356 &self,
357 record: KeyRecord,
358 max_active: NonZeroUsize,
359 now: Timestamp,
360 ) -> Result<(), KeyError> {
361 self.insert_key_within_audited(record, max_active, now)
362 .await
363 .map(|receipt| receipt.outcome)
364 }
365
366 async fn insert_key_within_audited(
367 &self,
368 record: KeyRecord,
369 max_active: NonZeroUsize,
370 now: Timestamp,
371 ) -> Result<AdminReceipt<()>, KeyError> {
372 self.insert_credential(record, Some((max_active, now)))
373 .await
374 }
375}
376
377impl PostgresStore {
378 async fn insert_credential(
383 &self,
384 record: KeyRecord,
385 bound: Option<(NonZeroUsize, Timestamp)>,
386 ) -> Result<AdminReceipt<()>, KeyError> {
387 let mut tx = self
388 .pool
389 .begin()
390 .await
391 .map_err(|e| KeyError::Storage(storage(e)))?;
392 let account = sqlx::query(ACCOUNT_LOCK_SQL)
400 .bind(id_bytes(record.account_id.0))
401 .fetch_optional(&mut *tx)
402 .await
403 .map_err(|e| KeyError::Storage(storage(e)))?;
404 if account.is_none() {
405 return Err(KeyError::UnknownAccount);
406 }
407
408 if let Some((max_active, now)) = bound {
409 let existing: Option<i32> = sqlx::query_scalar(
416 "SELECT 1 FROM tollgate_credential_keys WHERE key_id = $1 OR principal = $2",
417 )
418 .bind(id_bytes(record.key_id.0))
419 .bind(id_bytes(record.principal.0))
420 .fetch_optional(&mut *tx)
421 .await
422 .map_err(|e| KeyError::Storage(storage(e)))?;
423 if existing.is_some() {
424 return Err(KeyError::AlreadyExists);
425 }
426
427 let cutoff = StoredInstant::from(now);
428 let live: i64 = sqlx::query_scalar(LIVE_KEY_COUNT_SQL)
429 .bind(id_bytes(record.account_id.0))
430 .bind(cutoff.micros)
431 .bind(cutoff.submicro_nanos)
432 .fetch_one(&mut *tx)
433 .await
434 .map_err(|e| KeyError::Storage(storage(e)))?;
435 if u128::from(live.max(0).unsigned_abs())
436 >= u128::try_from(max_active.get()).unwrap_or(u128::MAX)
437 {
438 return Err(KeyError::ActiveKeyLimit { limit: max_active });
439 }
440 }
441
442 let expiry = record.not_after.map(StoredInstant::from);
443 let result = sqlx::query(
444 "INSERT INTO tollgate_credential_keys
445 (key_id, account_id, principal, digest, not_after_floor_us,
446 not_after_submicro_ns, not_after_is_lower_bound, revoked_at_us)
447 VALUES ($1, $2, $3, $4, $5, $6, FALSE, NULL)
448 ON CONFLICT (key_id) DO NOTHING",
449 )
450 .bind(id_bytes(record.key_id.0))
451 .bind(id_bytes(record.account_id.0))
452 .bind(id_bytes(record.principal.0))
453 .bind(record.digest.to_vec())
454 .bind(expiry.map(|expiry| expiry.micros))
455 .bind(expiry.map(|expiry| expiry.submicro_nanos))
456 .execute(&mut *tx)
457 .await;
458 let outcome = match result {
459 Ok(done) if done.rows_affected() == 0 => Err(KeyError::AlreadyExists),
460 Ok(_) => Ok(()),
461 Err(sqlx::Error::Database(e)) if e.is_unique_violation() => {
462 Err(KeyError::AlreadyExists)
463 }
464 Err(e) => Err(KeyError::Storage(storage(e))),
474 };
475 outcome?;
476 tx.commit()
477 .await
478 .map_err(|e| KeyError::Storage(storage(e)))?;
479 Ok(AdminReceipt::new(
480 (),
481 AdminState::Absent,
482 AdminState::Credential {
483 account_id: record.account_id,
484 key_id: record.key_id,
485 revoked: false,
486 },
487 ))
488 }
489}
490
491fn summary_from_row(row: sqlx::postgres::PgRow) -> Result<KeySummary, StoreError> {
492 let bytes: Vec<u8> = row.get(0);
493 let fixed: [u8; 16] = bytes
494 .as_slice()
495 .try_into()
496 .map_err(|_| StoreError("credential identifier is not 16 bytes".into()))?;
497 let not_after = match (row.get::<Option<i64>, _>(1), row.get::<Option<i16>, _>(2)) {
498 (None, None) if !row.get::<bool, _>(3) => None,
499 (Some(micros), Some(submicro_nanos)) => Some(
500 StoredInstant {
501 micros,
502 submicro_nanos,
503 }
504 .timestamp()?,
505 ),
506 _ => return Err(StoreError("incomplete stored credential expiry".into())),
507 };
508 Ok(KeySummary {
509 key_id: KeyId(u128::from_be_bytes(fixed)),
510 not_after,
511 revoked_at: row
512 .get::<Option<i64>, _>(4)
513 .map(|micros| {
514 StoredInstant {
515 micros,
516 submicro_nanos: 0,
517 }
518 .timestamp()
519 })
520 .transpose()?,
521 })
522}
523
524fn credential_from_row(row: sqlx::postgres::PgRow) -> Result<KeyRecord, StoreError> {
525 let id = |index| -> Result<u128, StoreError> {
526 let bytes: Vec<u8> = row.get(index);
527 let fixed: [u8; 16] = bytes
528 .as_slice()
529 .try_into()
530 .map_err(|_| StoreError("credential identifier is not 16 bytes".into()))?;
531 Ok(u128::from_be_bytes(fixed))
532 };
533 let key_id = KeyId(id(0)?);
534 let not_after = match (row.get::<Option<i64>, _>(4), row.get::<Option<i16>, _>(5)) {
535 (None, None) if !row.get::<bool, _>(6) => None,
536 (Some(micros), Some(submicro_nanos)) => Some(
537 StoredInstant {
538 micros,
539 submicro_nanos,
540 }
541 .timestamp()?,
542 ),
543 _ => return Err(StoreError("incomplete stored credential expiry".into())),
544 };
545 Ok(KeyRecord {
546 key_id,
547 account_id: AccountId(id(1)?),
548 principal: Principal(id(2)?),
549 digest: digest_from(row.get::<Vec<u8>, _>(3).as_slice(), key_id)?,
550 not_after,
551 })
552}
553
554#[async_trait]
555impl tollgate_store::KeySource for PostgresStore {
556 async fn active_keys_page(
557 &self,
558 now: Timestamp,
559 after: Option<KeyId>,
560 limit: NonZeroUsize,
561 ) -> Result<tollgate_store::KeyPage, StoreError> {
562 validate_key_page_limit(limit)?;
563 let cutoff = StoredInstant::from(now);
564 let mut tx = self.pool.begin().await.map_err(storage)?;
567 sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY")
568 .execute(&mut *tx)
569 .await
570 .map_err(storage)?;
571 let revision: i64 =
572 sqlx::query_scalar("SELECT revision FROM tollgate_credential_revision WHERE singleton")
573 .fetch_one(&mut *tx)
574 .await
575 .map_err(storage)?;
576 let revision = u64::try_from(revision)
577 .map_err(|_| StoreError("credential revision is negative".into()))?;
578 let sql = if after.is_some() {
580 "SELECT key_id, account_id, principal, digest, not_after_floor_us, not_after_submicro_ns, not_after_is_lower_bound
581 FROM tollgate_credential_keys WHERE revoked_at_us IS NULL
582 AND (not_after_floor_us IS NULL OR not_after_submicro_ns IS NULL
583 OR (not_after_floor_us, not_after_submicro_ns) > ($1, $2)) AND key_id > $3
584 ORDER BY key_id LIMIT $4"
585 } else {
586 "SELECT key_id, account_id, principal, digest, not_after_floor_us, not_after_submicro_ns, not_after_is_lower_bound
587 FROM tollgate_credential_keys WHERE revoked_at_us IS NULL
588 AND (not_after_floor_us IS NULL OR not_after_submicro_ns IS NULL
589 OR (not_after_floor_us, not_after_submicro_ns) > ($1, $2)) AND key_id >= $3
590 ORDER BY key_id LIMIT $4"
591 };
592 let rows = sqlx::query(sql)
593 .bind(cutoff.micros)
594 .bind(cutoff.submicro_nanos)
595 .bind(id_bytes(after.unwrap_or(KeyId(0)).0))
596 .bind((limit.get() + 1) as i64)
597 .fetch_all(&mut *tx)
598 .await
599 .map_err(storage)?;
600 tx.commit().await.map_err(storage)?;
601 let mut records = rows
604 .into_iter()
605 .map(credential_from_row)
606 .map(|row| row.map(tollgate_store::CredentialRecord::from))
607 .collect::<Result<Vec<_>, _>>()?;
608 let next_after = if records.len() > limit.get() {
609 records.pop();
610 records.last().map(|record| record.key_id)
611 } else {
612 None
613 };
614 tollgate_store::KeyPage::try_new(revision, now, after, limit, records, next_after)
615 }
616}
617
618#[cfg(test)]
619mod stored_id_tests {
620 use super::StoredId;
621
622 #[test]
623 fn malformed_storage_id_explains_both_accepted_representations() {
624 let error = serde_json::from_value::<StoredId>(serde_json::Value::Bool(true)).unwrap_err();
625 assert!(
626 error
627 .to_string()
628 .contains("a legacy u64 number or a canonical 128-bit identifier string"),
629 "unexpected diagnostic: {error}"
630 );
631 }
632}
633
634#[derive(Serialize)]
650struct StoredSnapshotRef<'a> {
651 account_id: StoredId,
652 key_id: Option<StoredId>,
653 status: &'a AccountStatus,
654 capacity_class: &'a CapacityClass,
657 enforcement_mode: &'a EnforcementMode,
658 valid_until: &'a Timestamp,
659 permissions: &'a PermissionBits,
660 limits: &'a ResolvedLimits,
661 cost_table: &'a Arc<CostTable>,
662 budget: Option<&'a BudgetView>,
666 policy_revision: &'a PolicyRevision,
671}
672
673impl<'a> From<&'a AccountSnapshot> for StoredSnapshotRef<'a> {
674 fn from(snapshot: &'a AccountSnapshot) -> Self {
675 StoredSnapshotRef {
676 account_id: StoredId(snapshot.account_id.0),
677 key_id: snapshot.key_id.map(|id| StoredId(id.0)),
678 status: &snapshot.status,
679 capacity_class: &snapshot.capacity_class,
680 enforcement_mode: &snapshot.enforcement_mode,
681 valid_until: &snapshot.valid_until,
682 permissions: &snapshot.permissions,
683 limits: &snapshot.limits,
684 cost_table: &snapshot.cost_table,
685 budget: snapshot.budget.as_ref(),
686 policy_revision: &snapshot.policy_revision,
687 }
688 }
689}
690
691#[derive(Deserialize)]
695struct StoredSnapshot {
696 account_id: StoredId,
697 key_id: Option<StoredId>,
698 status: AccountStatus,
699 #[serde(default)]
705 capacity_class: CapacityClass,
706 #[serde(default)]
717 enforcement_mode: EnforcementMode,
718 valid_until: Timestamp,
719 permissions: PermissionBits,
720 limits: ResolvedLimits,
721 cost_table: Arc<CostTable>,
722 #[serde(default)]
728 budget: Option<BudgetView>,
729 #[serde(default)]
735 policy_revision: PolicyRevision,
736}
737
738impl StoredSnapshot {
739 fn into_snapshot(self, generation: Generation) -> AccountSnapshot {
745 let builder = AccountSnapshot::builder(
746 AccountId(self.account_id.0),
747 generation,
748 self.status,
749 self.valid_until,
750 self.permissions,
751 self.limits,
752 self.cost_table,
753 )
754 .enforcement_mode(self.enforcement_mode)
755 .capacity_class(self.capacity_class)
756 .policy_revision(self.policy_revision);
757 match self.key_id {
758 Some(key_id) => builder.key_id(tollgate_core::KeyId(key_id.0)).build(),
759 None => builder.build(),
760 }
761 }
762}
763
764const ACTIVE_LEASE_SUM_SQL: &str =
782 "SELECT COALESCE(SUM(granted), 0)::BIGINT, COALESCE(SUM(used), 0)::BIGINT
783 FROM tollgate_leases WHERE account_id = $1 AND state = 0";
784
785pub(crate) const ACCOUNT_LOCK_SQL: &str =
804 "SELECT 1 FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE";
805
806pub(crate) const LIVE_KEY_COUNT_SQL: &str = "SELECT count(*) FROM tollgate_credential_keys
809 WHERE account_id = $1
810 AND revoked_at_us IS NULL
811 AND (not_after_floor_us IS NULL OR not_after_submicro_ns IS NULL
812 OR (not_after_floor_us, not_after_submicro_ns) > ($2, $3))";
813
814const RECLAIM_DUE_LEASES_SQL: &str = "SELECT lease_id, account_id, granted, used
815 FROM tollgate_leases
816 WHERE state = 0 AND (expires_at_floor_us, expires_at_submicro_ns) <= ($1, $2)
817 ORDER BY expires_at_floor_us, expires_at_submicro_ns
818 LIMIT $3 FOR UPDATE SKIP LOCKED";
819
820const DUE_PERIODS_SQL: &str = "WITH due AS (
849 SELECT account_id, allowance_balance AS prior, budget_allowance AS allowance,
850 period_start_us AS crossed_from
851 FROM tollgate_accounts
852 WHERE budget_allowance IS NOT NULL
853 AND budget_period = $1 AND period_start_us < $2
854 ORDER BY period_start_us LIMIT $3 FOR UPDATE SKIP LOCKED
855 )
856 UPDATE tollgate_accounts AS account SET
857 deposited = account.deposited + due.allowance,
858 expired = account.expired + due.prior,
859 balance = account.balance - due.prior + due.allowance,
860 allowance_balance = due.allowance,
861 period_start_us = $2
862 FROM due
863 WHERE account.account_id = due.account_id
864 RETURNING due.account_id, due.allowance, due.prior, due.crossed_from";
865
866fn id_bytes(id: u128) -> Vec<u8> {
867 id.to_be_bytes().to_vec()
868}
869
870fn id_from(bytes: &[u8]) -> u128 {
871 let mut buf = [0u8; 16];
872 buf.copy_from_slice(bytes);
873 u128::from_be_bytes(buf)
874}
875
876fn to_i64(units: CostUnits, what: &str) -> Result<i64, StoreError> {
877 i64::try_from(units.get()).map_err(|_| StoreError(format!("{what} exceeds i64 range")))
878}
879
880fn stored_fence(value: i64) -> Result<FencingToken, StoreError> {
883 u64::try_from(value)
884 .ok()
885 .filter(|value| *value > 0)
886 .map(FencingToken)
887 .ok_or_else(|| StoreError(format!("stored fencing token is not positive: {value}")))
888}
889
890fn to_units(value: i64, what: &str) -> Result<CostUnits, StoreError> {
894 u64::try_from(value)
895 .map(CostUnits)
896 .map_err(|_| StoreError(format!("{what} is negative in storage: {value}")))
897}
898
899fn ts_micros(ts: Timestamp) -> i64 {
900 ts.as_microsecond()
901}
902
903fn micros_ts(value: i64, what: &str) -> Result<Timestamp, StoreError> {
906 tollgate_store::clock::timestamp_from_micros(value).map_err(|e| {
907 StoreError(format!(
908 "{what} is not a representable instant: {value} ({e})"
909 ))
910 })
911}
912
913fn storage(e: sqlx::Error) -> StoreError {
914 StoreError(format!("postgres: {e}"))
915}
916
917fn alloc_storage(e: sqlx::Error) -> AllocateError {
918 AllocateError::Storage(storage(e))
919}
920
921async fn finish_transaction<T, E>(
934 tx: Transaction<'_, Postgres>,
935 result: Result<T, E>,
936) -> Result<T, E>
937where
938 E: From<StoreError> + std::fmt::Display,
939{
940 match result {
941 Ok(value) => {
942 tx.commit().await.map_err(|e| E::from(storage(e)))?;
943 Ok(value)
944 }
945 Err(error) => {
946 if let Err(rollback_error) = tx.rollback().await {
947 return Err(E::from(StoreError(format!(
948 "operation failed ({error}); transaction rollback failed ({})",
949 storage(rollback_error)
950 ))));
951 }
952 Err(error)
953 }
954 }
955}
956
957pub struct PostgresStore {
994 pool: PgPool,
995 policy: GrantPolicy,
996 push: broadcast::Sender<SnapshotPush>,
997}
998
999#[derive(Debug, Clone, Copy)]
1004pub struct PoolConfig {
1005 pub max_connections: u32,
1008 pub acquire_timeout: std::time::Duration,
1011}
1012
1013impl Default for PoolConfig {
1014 fn default() -> Self {
1015 PoolConfig {
1016 max_connections: 16,
1017 acquire_timeout: std::time::Duration::from_secs(5),
1018 }
1019 }
1020}
1021
1022impl PoolConfig {
1023 pub fn validate(&self) -> Result<(), StoreError> {
1029 if self.max_connections == 0 {
1030 return Err(StoreError("max_connections must be positive".into()));
1031 }
1032 if self.acquire_timeout.is_zero() {
1033 return Err(StoreError("acquire_timeout must be positive".into()));
1034 }
1035 Ok(())
1036 }
1037}
1038
1039impl PostgresStore {
1040 pub async fn connect(url: &str, policy: GrantPolicy) -> Result<Arc<Self>, StoreError> {
1042 Self::connect_with(url, policy, PoolConfig::default()).await
1043 }
1044
1045 pub async fn connect_with(
1054 url: &str,
1055 policy: GrantPolicy,
1056 pool_config: PoolConfig,
1057 ) -> Result<Arc<Self>, StoreError> {
1058 policy
1059 .validate()
1060 .map_err(|e| StoreError(format!("invalid grant policy: {e}")))?;
1061 pool_config.validate()?;
1062 let pool = PgPoolOptions::new()
1063 .max_connections(pool_config.max_connections)
1064 .acquire_timeout(pool_config.acquire_timeout)
1065 .connect(url)
1066 .await
1067 .map_err(storage)?;
1068 sqlx::migrate!("./migrations")
1069 .run(&pool)
1070 .await
1071 .map_err(|e| StoreError(format!("migrate: {e}")))?;
1072 let (push, _) = broadcast::channel(PUSH_CHANNEL_CAPACITY);
1073 Ok(Arc::new(PostgresStore { pool, policy, push }))
1074 }
1075
1076 fn push_to_subscribers(&self, push: SnapshotPush) {
1082 let principal = push.principal;
1083 let subscribers = self.push.send(push).unwrap_or(0);
1084 tracing::debug!(
1085 %principal,
1086 subscribers,
1087 "snapshot pushed to subscribers"
1088 );
1089 }
1090
1091 fn announce_removal(&self, principal: Principal, receipt: &AdminReceipt<()>) {
1093 if receipt.before != receipt.after
1094 && let AdminState::Snapshot { generation, .. } = receipt.after
1095 {
1096 self.push_to_subscribers(SnapshotPush {
1097 principal,
1098 resolution: SnapshotResolution::Revoked { generation },
1099 });
1100 }
1101 }
1102
1103 pub async fn balance(&self, account: AccountId) -> Result<CostUnits, StoreError> {
1113 let row = sqlx::query("SELECT balance FROM tollgate_accounts WHERE account_id = $1")
1114 .bind(id_bytes(account.0))
1115 .fetch_optional(&self.pool)
1116 .await
1117 .map_err(storage)?;
1118 row.map(|r| to_units(r.get::<i64, _>(0), "account balance"))
1119 .transpose()
1120 .map(|units| units.unwrap_or(CostUnits::ZERO))
1121 }
1122
1123 pub async fn usage_recorded(&self, account: AccountId) -> Result<CostUnits, StoreError> {
1130 let row = sqlx::query("SELECT usage_recorded FROM tollgate_accounts WHERE account_id = $1")
1131 .bind(id_bytes(account.0))
1132 .fetch_optional(&self.pool)
1133 .await
1134 .map_err(storage)?;
1135 row.map(|r| to_units(r.get::<i64, _>(0), "account usage_recorded"))
1136 .transpose()
1137 .map(|units| units.unwrap_or(CostUnits::ZERO))
1138 }
1139
1140 pub async fn conservation(
1155 &self,
1156 account: AccountId,
1157 ) -> Result<Option<Conservation>, StoreError> {
1158 let mut tx = self.pool.begin().await.map_err(storage)?;
1172 let result = async {
1173 sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY")
1174 .execute(&mut *tx)
1175 .await
1176 .map_err(storage)?;
1177 let account_row = sqlx::query(
1178 "SELECT deposited, balance, usage_recorded, settlement_loss, overage_recorded,
1179 expired
1180 FROM tollgate_accounts WHERE account_id = $1",
1181 )
1182 .bind(id_bytes(account.0))
1183 .fetch_optional(&mut *tx)
1184 .await
1185 .map_err(storage)?;
1186 let lease_row = sqlx::query(ACTIVE_LEASE_SUM_SQL)
1187 .bind(id_bytes(account.0))
1188 .fetch_one(&mut *tx)
1189 .await
1190 .map_err(storage)?;
1191 Ok::<_, StoreError>((account_row, lease_row))
1192 }
1193 .await;
1194 let (account_row, lease_row) = finish_transaction(tx, result).await?;
1195
1196 let Some(row) = account_row else {
1197 return Ok(None);
1198 };
1199 let active_grants = to_units(lease_row.get::<i64, _>(0), "active lease grants")?;
1200 let active_used = to_units(lease_row.get::<i64, _>(1), "active lease usage")?;
1201 let recorded = to_units(row.get::<i64, _>(2), "usage_recorded")?;
1202 Ok(Some(Conservation {
1203 deposited: to_units(row.get::<i64, _>(0), "deposited")?,
1204 overage_recorded: to_units(row.get::<i64, _>(4), "overage_recorded")?,
1205 balance: to_units(row.get::<i64, _>(1), "balance")?,
1206 active_lease_grants: active_grants,
1207 settled_usage: recorded.checked_sub(active_used).ok_or_else(|| {
1215 StoreError(format!(
1216 "active lease usage {} exceeds recorded usage {} for account {account}",
1217 active_used.get(),
1218 recorded.get()
1219 ))
1220 })?,
1221 settlement_loss: to_units(row.get::<i64, _>(3), "settlement_loss")?,
1222 expired: to_units(row.get::<i64, _>(5), "expired")?,
1223 }))
1224 }
1225}
1226
1227struct LockedLeaseRow {
1231 account_id: Vec<u8>,
1232 fencing_token: i64,
1233 granted: i64,
1234 used: i64,
1235 credited: i64,
1236 expires_at: StoredInstant,
1237 state: i16,
1238 from_allowance: i64,
1243 period_start_us: i64,
1244}
1245
1246struct ReleasedCredit {
1249 account: AccountId,
1250 restored: CostUnits,
1251 preserves_funding: bool,
1252}
1253
1254async fn lock_lease(
1256 tx: &mut Transaction<'_, Postgres>,
1257 lease_id: LeaseId,
1258) -> Result<Option<LockedLeaseRow>, sqlx::Error> {
1259 let row = sqlx::query(
1260 "SELECT account_id, fencing_token, granted, used, credited, expires_at_floor_us, state,
1261 from_allowance, period_start_us, expires_at_submicro_ns
1262 FROM tollgate_leases WHERE lease_id = $1 FOR UPDATE",
1263 )
1264 .bind(id_bytes(lease_id.0))
1265 .fetch_optional(&mut **tx)
1266 .await?;
1267 Ok(row.map(|row| LockedLeaseRow {
1268 account_id: row.get(0),
1269 fencing_token: row.get(1),
1270 granted: row.get(2),
1271 used: row.get(3),
1272 credited: row.get(4),
1273 expires_at: StoredInstant {
1274 micros: row.get(5),
1275 submicro_nanos: row.get(9),
1276 },
1277 state: row.get(6),
1278 from_allowance: row.get(7),
1279 period_start_us: row.get(8),
1280 }))
1281}
1282
1283struct Exchange {
1286 floor: CostUnits,
1290 needed: CostUnits,
1293 preserves_funding: bool,
1296}
1297
1298impl Exchange {
1299 const ACQUIRE: Exchange = Exchange {
1300 floor: CostUnits::ZERO,
1301 needed: CostUnits::ZERO,
1302 preserves_funding: true,
1303 };
1304}
1305
1306impl PostgresStore {
1307 async fn acquire_in_tx(
1309 &self,
1310 tx: &mut Transaction<'_, Postgres>,
1311 account: AccountId,
1312 requested: CostUnits,
1313 expires_at: Timestamp,
1314 exchange: Exchange,
1315 ) -> Result<Allocation, AllocateError> {
1316 let Exchange {
1317 floor,
1318 needed,
1319 preserves_funding: settlement_preserves_funding,
1320 } = exchange;
1321 let row = sqlx::query(
1322 "SELECT balance, status, next_fence, allowance_balance, period_start_us
1323 FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE",
1324 )
1325 .bind(id_bytes(account.0))
1326 .fetch_optional(&mut **tx)
1327 .await
1328 .map_err(alloc_storage)?
1329 .ok_or(AllocateError::UnknownAccount)?;
1330
1331 if decode_status(row.get::<String, _>(1)).map_err(AllocateError::Storage)?
1335 != AccountStatus::Active
1336 {
1337 return Err(AllocateError::AccountInactive);
1338 }
1339 let balance =
1340 to_units(row.get::<i64, _>(0), "account balance").map_err(AllocateError::Storage)?;
1341 let granted = match self
1347 .policy
1348 .consolidation_grant(requested, balance, floor, needed)
1349 {
1350 Some(granted) => granted,
1351 None => {
1352 if settlement_preserves_funding && !requested.is_zero() {
1355 let budget = sqlx::query(BUDGET_VIEW_SQL)
1356 .bind(id_bytes(account.0))
1357 .fetch_one(&mut **tx)
1358 .await
1359 .map_err(alloc_storage)?;
1360 let evidence = budget_view(&budget)
1361 .map_err(AllocateError::Storage)?
1362 .shortfall();
1363 return Err(match evidence.exhaustion() {
1364 Some(exhausted) => AllocateError::BalanceExhausted(exhausted),
1365 None => AllocateError::BalanceInsufficient(evidence),
1366 });
1367 }
1368 return Err(AllocateError::InsufficientBalance);
1369 }
1370 };
1371 let granted_i = to_i64(granted, "grant").map_err(AllocateError::Storage)?;
1376 let allowance_balance = row.get::<i64, _>(3);
1377 let from_allowance = granted_i.min(allowance_balance);
1378 let period_start_us = row.get::<i64, _>(4);
1379 let fence = row.get::<i64, _>(2);
1380 let fence_token = stored_fence(fence).map_err(AllocateError::Storage)?;
1381
1382 let lease_id = LeaseId(uuid::Uuid::new_v4().as_u128());
1383
1384 let budget = sqlx::query(GRANT_DEBIT_SQL)
1388 .bind(id_bytes(account.0))
1389 .bind(granted_i)
1390 .bind(from_allowance)
1391 .fetch_one(&mut **tx)
1392 .await
1393 .map_err(alloc_storage)?;
1394 let funding = budget_view(&budget)
1395 .map_err(AllocateError::Storage)?
1396 .shortfall();
1397 sqlx::query(
1398 "INSERT INTO tollgate_leases
1399 (lease_id, account_id, fencing_token, granted, used, credited, expires_at_floor_us,
1400 state, from_allowance, period_start_us, expires_at_submicro_ns, expiry_is_upper_bound)
1401 VALUES ($1, $2, $3, $4, 0, 0, $5, 0, $6, $7, $8, FALSE)",
1402 )
1403 .bind(id_bytes(lease_id.0))
1404 .bind(id_bytes(account.0))
1405 .bind(fence)
1406 .bind(granted_i)
1407 .bind(StoredInstant::from(expires_at).micros)
1408 .bind(from_allowance)
1409 .bind(period_start_us)
1410 .bind(StoredInstant::from(expires_at).submicro_nanos)
1411 .execute(&mut **tx)
1412 .await
1413 .map_err(alloc_storage)?;
1414
1415 Ok(Allocation {
1416 grant: LeaseGrant {
1417 lease_id,
1418 account_id: account,
1419 fencing_token: fence_token,
1420 units: granted,
1421 expires_at,
1422 },
1423 funding: Some(funding),
1424 })
1425 }
1426
1427 async fn release_in_tx(
1430 &self,
1431 tx: &mut Transaction<'_, Postgres>,
1432 lease_id: LeaseId,
1433 fencing_token: FencingToken,
1434 unspent: CostUnits,
1435 now: Timestamp,
1436 ) -> Result<ReleasedCredit, AllocateError> {
1437 let LockedLeaseRow {
1438 account_id,
1439 fencing_token: fence,
1440 granted,
1441 used,
1442 credited: _credited,
1443 expires_at,
1444 state,
1445 from_allowance,
1446 period_start_us,
1447 } = lock_lease(tx, lease_id)
1448 .await
1449 .map_err(alloc_storage)?
1450 .ok_or(AllocateError::UnknownLease)?;
1451
1452 if stored_fence(fence).map_err(AllocateError::Storage)? != fencing_token {
1453 return Err(AllocateError::Fenced);
1454 }
1455 let expires_at = expires_at.timestamp().map_err(AllocateError::Storage)?;
1458 if state != STATE_ACTIVE
1459 || self
1460 .policy
1461 .reclaim_cutoff(now)
1462 .is_some_and(|cutoff| expires_at <= cutoff)
1463 {
1464 return Err(AllocateError::LeaseNotActive);
1465 }
1466 let account = AccountId(id_from(&account_id));
1470 let unspent_i = to_i64(unspent, "unspent").map_err(AllocateError::Storage)?;
1471 let loss = granted
1472 .checked_sub(
1473 used.checked_add(unspent_i)
1474 .ok_or(AllocateError::InvalidRelease)?,
1475 )
1476 .filter(|l| *l >= 0)
1477 .ok_or(AllocateError::InvalidRelease)?;
1478
1479 sqlx::query("UPDATE tollgate_leases SET state = $2, credited = $3 WHERE lease_id = $1")
1480 .bind(id_bytes(lease_id.0))
1481 .bind(STATE_RELEASED)
1482 .bind(unspent_i)
1483 .execute(&mut **tx)
1484 .await
1485 .map_err(alloc_storage)?;
1486 let from_topup = granted
1492 .checked_sub(from_allowance)
1493 .filter(|t| *t >= 0)
1494 .ok_or_else(|| {
1495 AllocateError::Storage(StoreError(format!(
1496 "lease allowance funding {from_allowance} exceeds its grant {granted}"
1497 )))
1498 })?;
1499 let to_topup = from_topup.min(unspent_i);
1500 let to_allowance = unspent_i - to_topup;
1501
1502 let restored: i64 = sqlx::query_scalar(
1509 "UPDATE tollgate_accounts SET
1510 balance = balance + $2
1511 + CASE WHEN period_start_us > $5 THEN 0 ELSE $3 END,
1512 allowance_balance = allowance_balance
1513 + CASE WHEN period_start_us > $5 THEN 0 ELSE $3 END,
1514 expired = expired + CASE WHEN period_start_us > $5 THEN $3 ELSE 0 END,
1515 settlement_loss = settlement_loss + $4
1516 WHERE account_id = $1
1517 RETURNING $2 + CASE WHEN period_start_us > $5 THEN 0 ELSE $3 END",
1518 )
1519 .bind(account_id)
1520 .bind(to_topup)
1521 .bind(to_allowance)
1522 .bind(loss)
1523 .bind(period_start_us)
1524 .fetch_one(&mut **tx)
1525 .await
1526 .map_err(alloc_storage)?;
1527 Ok(ReleasedCredit {
1528 account,
1529 preserves_funding: loss == 0 && restored == unspent_i,
1530 restored: to_units(restored, "restored release credit")
1531 .map_err(AllocateError::Storage)?,
1532 })
1533 }
1534
1535 fn grant_expiry(
1539 &self,
1540 ttl: SignedDuration,
1541 now: Timestamp,
1542 ) -> Result<Timestamp, AllocateError> {
1543 if ttl <= SignedDuration::ZERO {
1544 return Err(AllocateError::InvalidTtl);
1545 }
1546 now.checked_add(ttl.min(self.policy.max_ttl))
1547 .map_err(|e| AllocateError::Storage(StoreError(format!("ttl overflow: {e}"))))
1548 }
1549}
1550
1551#[async_trait]
1552impl LeaseAllocator for PostgresStore {
1553 async fn acquire(
1554 &self,
1555 account: AccountId,
1556 requested: CostUnits,
1557 ttl: SignedDuration,
1558 now: Timestamp,
1559 ) -> Result<Allocation, AllocateError> {
1560 let expires_at = self.grant_expiry(ttl, now)?;
1561 let mut tx = self.pool.begin().await.map_err(alloc_storage)?;
1562 let result = self
1563 .acquire_in_tx(&mut tx, account, requested, expires_at, Exchange::ACQUIRE)
1564 .await;
1565 finish_transaction(tx, result).await
1566 }
1567
1568 async fn release(
1569 &self,
1570 lease_id: LeaseId,
1571 fencing_token: FencingToken,
1572 unspent: CostUnits,
1573 now: Timestamp,
1574 ) -> Result<(), AllocateError> {
1575 let mut tx = self.pool.begin().await.map_err(alloc_storage)?;
1576 let result = self
1577 .release_in_tx(&mut tx, lease_id, fencing_token, unspent, now)
1578 .await
1579 .map(|_| ());
1580 finish_transaction(tx, result).await
1581 }
1582
1583 async fn consolidate(
1584 &self,
1585 lease_id: LeaseId,
1586 fencing_token: FencingToken,
1587 unspent: CostUnits,
1588 requested: CostUnits,
1589 needed: CostUnits,
1590 ttl: SignedDuration,
1591 now: Timestamp,
1592 ) -> Result<Allocation, AllocateError> {
1593 let expires_at = self.grant_expiry(ttl, now)?;
1594 let mut tx = self.pool.begin().await.map_err(alloc_storage)?;
1595 let result = async {
1604 let released = self
1605 .release_in_tx(&mut tx, lease_id, fencing_token, unspent, now)
1606 .await?;
1607 self.acquire_in_tx(
1608 &mut tx,
1609 released.account,
1610 requested,
1611 expires_at,
1612 Exchange {
1613 floor: released.restored,
1614 needed,
1615 preserves_funding: released.preserves_funding,
1616 },
1617 )
1618 .await
1619 }
1620 .await;
1621 finish_transaction(tx, result).await
1622 }
1623
1624 async fn reclaim_expired_batch(
1625 &self,
1626 now: Timestamp,
1627 limit: NonZeroUsize,
1628 ) -> Result<ReclaimBatch, StoreError> {
1629 let limit_i = i64::try_from(limit.get())
1630 .map_err(|_| StoreError(format!("reclaim batch limit exceeds i64 range: {limit}")))?;
1631 let Some(cutoff) = self.policy.reclaim_cutoff(now) else {
1632 return ReclaimBatch::try_new(Vec::new(), limit);
1633 };
1634 let cutoff = StoredInstant::from(cutoff);
1635 let mut tx = self.pool.begin().await.map_err(storage)?;
1636 let result = async {
1637 let rows = sqlx::query(RECLAIM_DUE_LEASES_SQL)
1643 .bind(cutoff.micros)
1644 .bind(cutoff.submicro_nanos)
1645 .bind(limit_i)
1646 .fetch_all(&mut *tx)
1647 .await
1648 .map_err(storage)?;
1649
1650 let mut reclaimed = Vec::with_capacity(rows.len());
1655 let mut lease_ids = Vec::with_capacity(rows.len());
1656 let mut forfeits: std::collections::BTreeMap<Vec<u8>, i64> =
1657 std::collections::BTreeMap::new();
1658 for row in rows {
1659 let lease_bytes: Vec<u8> = row.get(0);
1660 let account_bytes: Vec<u8> = row.get(1);
1661 let granted = row.get::<i64, _>(2);
1662 let used = row.get::<i64, _>(3);
1663 let forfeited = granted.checked_sub(used).ok_or_else(|| {
1667 StoreError(format!(
1668 "reclaim remainder overflow: granted {granted}, used {used}"
1669 ))
1670 })?;
1671 let forfeited_units = to_units(forfeited, "reclaim remainder")?;
1672 let total = forfeits.entry(account_bytes.clone()).or_default();
1673 *total = total
1674 .checked_add(forfeited)
1675 .ok_or_else(|| StoreError("reclaim loss sum overflow".into()))?;
1676 lease_ids.push(lease_bytes.clone());
1677 reclaimed.push(ReclaimedLease {
1678 lease_id: LeaseId(id_from(&lease_bytes)),
1679 account_id: AccountId(id_from(&account_bytes)),
1680 forfeited: forfeited_units,
1681 });
1682 }
1683
1684 let batch = ReclaimBatch::try_new(reclaimed, limit)?;
1685 if batch.is_empty() {
1686 return Ok(batch);
1687 }
1688
1689 let (account_ids, account_forfeits): (Vec<_>, Vec<_>) = forfeits.into_iter().unzip();
1690 let expected_lease_rows = u64::try_from(lease_ids.len())
1691 .map_err(|_| StoreError("reclaim lease row count exceeds u64 range".into()))?;
1692 let expected_account_rows = u64::try_from(account_ids.len())
1693 .map_err(|_| StoreError("reclaim account row count exceeds u64 range".into()))?;
1694
1695 let locked_accounts = sqlx::query(
1701 "SELECT account_id FROM tollgate_accounts
1702 WHERE account_id = ANY($1) ORDER BY account_id FOR UPDATE",
1703 )
1704 .bind(&account_ids)
1705 .fetch_all(&mut *tx)
1706 .await
1707 .map_err(storage)?;
1708 if locked_accounts.len() != account_ids.len() {
1709 return Err(StoreError(format!(
1710 "reclaim locked {} of {} referenced account rows",
1711 locked_accounts.len(),
1712 account_ids.len()
1713 )));
1714 }
1715
1716 let updated_leases = sqlx::query(
1717 "UPDATE tollgate_leases
1718 SET state = $2, credited = 0
1719 WHERE lease_id = ANY($1) AND state = $3",
1720 )
1721 .bind(&lease_ids)
1722 .bind(STATE_EXPIRED)
1723 .bind(STATE_ACTIVE)
1724 .execute(&mut *tx)
1725 .await
1726 .map_err(storage)?;
1727 if updated_leases.rows_affected() != expected_lease_rows {
1728 return Err(StoreError(format!(
1729 "reclaim updated {} of {} locked lease rows",
1730 updated_leases.rows_affected(),
1731 lease_ids.len()
1732 )));
1733 }
1734
1735 let updated_accounts = sqlx::query(
1736 "UPDATE tollgate_accounts AS account
1737 SET settlement_loss = account.settlement_loss + delta.forfeited
1738 FROM UNNEST($1::bytea[], $2::bigint[]) AS delta(account_id, forfeited)
1739 WHERE account.account_id = delta.account_id",
1740 )
1741 .bind(&account_ids)
1742 .bind(&account_forfeits)
1743 .execute(&mut *tx)
1744 .await
1745 .map_err(storage)?;
1746 if updated_accounts.rows_affected() != expected_account_rows {
1747 return Err(StoreError(format!(
1748 "reclaim updated {} of {} locked account rows",
1749 updated_accounts.rows_affected(),
1750 account_ids.len()
1751 )));
1752 }
1753
1754 Ok(batch)
1755 }
1756 .await;
1757 finish_transaction(tx, result).await
1758 }
1759}
1760
1761#[async_trait]
1762impl UsageSink for PostgresStore {
1763 async fn ingest(
1764 &self,
1765 events: &[UsageEvent],
1766 _now: Timestamp,
1767 ) -> Result<IngestReport, IngestError> {
1768 let mut report = IngestReport {
1776 unattributed: Some(0),
1777 ..IngestReport::default()
1778 };
1779 if events.is_empty() {
1780 return Ok(report);
1781 }
1782
1783 struct PreparedEvent<'a> {
1789 event: &'a UsageEvent,
1790 request_id: Vec<u8>,
1791 account_id: Vec<u8>,
1792 key_id: Option<Vec<u8>>,
1793 lease_id: Option<Vec<u8>>,
1797 occurred_at_us: i64,
1798 }
1799 let prepared: Vec<PreparedEvent<'_>> = events
1800 .iter()
1801 .map(|event| {
1802 Ok(PreparedEvent {
1803 event,
1804 request_id: id_bytes(event.request_id.0),
1805 account_id: id_bytes(event.account_id.0),
1806 key_id: event.key_id.map(|id| id_bytes(id.0)),
1807 lease_id: event.source.lease_id().map(|id| id_bytes(id.0)),
1808 occurred_at_us: ts_micros(event.occurred_at),
1809 })
1810 })
1811 .collect::<Result<_, StoreError>>()?;
1812
1813 let mut tx = self.pool.begin().await.map_err(storage)?;
1814 let result: Result<_, IngestError> = async {
1815 let lease_ids: Vec<Vec<u8>> = prepared
1824 .iter()
1825 .filter_map(|event| event.lease_id.clone())
1826 .collect::<std::collections::BTreeSet<_>>()
1827 .into_iter()
1828 .collect();
1829 struct LeaseRow {
1830 account_id: Vec<u8>,
1831 fence: i64,
1832 granted: i64,
1833 used: i64,
1834 used_delta: Option<NonZeroI64>,
1835 credited: i64,
1836 settled: bool,
1837 }
1838 let rows = sqlx::query(
1839 "SELECT lease_id, account_id, fencing_token, granted, used, credited, state
1840 FROM tollgate_leases
1841 WHERE lease_id = ANY($1)
1842 ORDER BY account_id, lease_id FOR UPDATE",
1843 )
1844 .bind(&lease_ids)
1845 .fetch_all(&mut *tx)
1846 .await
1847 .map_err(storage)?;
1848 let mut leases: std::collections::BTreeMap<Vec<u8>, LeaseRow> =
1849 std::collections::BTreeMap::new();
1850 for row in rows {
1851 let lease_id: Vec<u8> = row.get(0);
1852 let fence: i64 = row.get(2);
1853 let granted: i64 = row.get(3);
1854 let used: i64 = row.get(4);
1855 let credited: i64 = row.get(5);
1856 leases.insert(
1857 lease_id,
1858 LeaseRow {
1859 account_id: row.get(1),
1860 fence,
1861 granted,
1862 used,
1863 used_delta: None,
1864 credited,
1865 settled: row.get::<i16, _>(6) != STATE_ACTIVE,
1866 },
1867 );
1868 }
1869
1870 let request_ids: Vec<Vec<u8>> = prepared
1872 .iter()
1873 .map(|event| event.request_id.clone())
1874 .collect();
1875 let mut seen: std::collections::HashSet<Vec<u8>> = sqlx::query(
1876 "SELECT request_id FROM tollgate_usage_events WHERE request_id = ANY($1)",
1877 )
1878 .bind(&request_ids)
1879 .fetch_all(&mut *tx)
1880 .await
1881 .map_err(storage)?
1882 .into_iter()
1883 .map(|row| row.get::<Vec<u8>, _>(0))
1884 .collect();
1885
1886 let overage_account_ids: Vec<Vec<u8>> = prepared
1902 .iter()
1903 .filter(|event| event.lease_id.is_none())
1904 .map(|event| event.account_id.clone())
1905 .collect::<std::collections::BTreeSet<_>>()
1906 .into_iter()
1907 .collect();
1908 let known_overage_accounts: std::collections::HashSet<Vec<u8>> =
1909 if overage_account_ids.is_empty() {
1910 std::collections::HashSet::new()
1911 } else {
1912 sqlx::query(
1913 "SELECT account_id FROM tollgate_accounts WHERE account_id = ANY($1)",
1914 )
1915 .bind(&overage_account_ids)
1916 .fetch_all(&mut *tx)
1917 .await
1918 .map_err(storage)?
1919 .into_iter()
1920 .map(|row| row.get::<Vec<u8>, _>(0))
1921 .collect()
1922 };
1923
1924 struct Accepted {
1928 event_index: usize,
1929 settled: bool,
1930 fence: Option<i64>,
1935 overage: bool,
1938 units: i64,
1941 }
1942 let mut accepted: Vec<Accepted> = Vec::with_capacity(prepared.len());
1943 for (event_index, event) in prepared.iter().enumerate() {
1944 if seen.contains(event.request_id.as_slice()) {
1945 report.duplicate += 1;
1946 continue;
1947 }
1948 let Some(lease_key) = event.lease_id.as_deref() else {
1949 if !known_overage_accounts.contains(event.account_id.as_slice()) {
1958 report.rejected += 1;
1959 continue;
1960 }
1961 let Ok(units) = i64::try_from(event.event.units.get()) else {
1962 report.rejected += 1;
1965 continue;
1966 };
1967 seen.insert(event.request_id.clone());
1968 accepted.push(Accepted {
1969 event_index,
1970 settled: false,
1974 fence: None,
1975 overage: true,
1976 units,
1977 });
1978 report.accepted += 1;
1979 continue;
1980 };
1981 let Some(lease) = leases.get_mut(lease_key) else {
1982 report.rejected += 1;
1983 continue;
1984 };
1985 if Some(stored_fence(lease.fence)?) != event.event.source.fencing_token()
1986 || lease.account_id.as_slice() != event.account_id.as_slice()
1987 {
1988 report.rejected += 1;
1989 continue;
1990 }
1991 to_units(lease.granted, "lease granted")?;
1992 to_units(lease.used, "lease used")?;
1993 to_units(lease.credited, "lease credited")?;
1994 let Ok(units) = i64::try_from(event.event.units.get()) else {
1995 report.rejected += 1;
1996 continue;
1997 };
1998 let committed = lease
1999 .used
2000 .checked_add(lease.used_delta.map_or(0, NonZeroI64::get))
2001 .and_then(|used| used.checked_add(lease.credited))
2002 .ok_or_else(|| {
2003 StoreError(format!(
2004 "lease accounting overflow for {:#034x}",
2005 id_from(lease_key)
2006 ))
2007 })?;
2008 let remaining = lease.granted.checked_sub(committed).ok_or_else(|| {
2009 StoreError(format!(
2010 "lease accounting exceeds grant for {:#034x}: granted {}, committed {committed}",
2011 id_from(lease_key), lease.granted
2012 ))
2013 })?;
2014 if units > remaining {
2015 report.rejected += 1;
2016 continue;
2017 }
2018 let used_delta = lease
2019 .used_delta
2020 .map_or(0, NonZeroI64::get)
2021 .checked_add(units)
2022 .ok_or_else(|| {
2023 StoreError(format!(
2024 "lease usage delta overflow for {:#034x}",
2025 id_from(lease_key)
2026 ))
2027 })?;
2028 lease.used_delta = NonZeroI64::new(used_delta);
2029 seen.insert(event.request_id.clone());
2030 accepted.push(Accepted {
2031 event_index,
2032 settled: lease.settled,
2033 fence: Some(lease.fence),
2034 overage: false,
2035 units,
2036 });
2037 report.accepted += 1;
2038 }
2039
2040 if accepted.is_empty() {
2041 return Ok(report);
2042 }
2043
2044 #[derive(Default)]
2045 struct AccountDelta {
2046 usage: i64,
2047 loss: i64,
2048 overage: i64,
2052 }
2053
2054 let (mut rid, mut acct, mut lease, mut fence, mut units, mut at, mut revision, mut keys) = (
2056 Vec::with_capacity(accepted.len()),
2057 Vec::with_capacity(accepted.len()),
2058 Vec::with_capacity(accepted.len()),
2059 Vec::with_capacity(accepted.len()),
2060 Vec::with_capacity(accepted.len()),
2061 Vec::with_capacity(accepted.len()),
2062 Vec::with_capacity(accepted.len()),
2063 Vec::with_capacity(accepted.len()),
2064 );
2065 let mut account_deltas: std::collections::BTreeMap<Vec<u8>, AccountDelta> =
2066 std::collections::BTreeMap::new();
2067 for accepted_event in &accepted {
2068 let event = &prepared[accepted_event.event_index];
2069 rid.push(event.request_id.clone());
2070 acct.push(event.account_id.clone());
2071 keys.push(event.key_id.clone());
2072 lease.push(event.lease_id.clone());
2073 fence.push(accepted_event.fence);
2074 units.push(accepted_event.units);
2075 at.push(event.occurred_at_us);
2076 revision.push(event.event.policy_revision.as_bytes().to_vec());
2081 if accepted_event.overage {
2082 debug_assert!(
2083 event.lease_id.is_none() && accepted_event.fence.is_none(),
2084 "an overage row must carry neither half of a capability"
2085 );
2086 }
2087
2088 let entry = account_deltas.entry(event.account_id.clone()).or_default();
2089 entry.usage = entry.usage.checked_add(accepted_event.units).ok_or_else(|| {
2090 IngestError::Refused(StoreError(format!(
2091 "usage delta overflow for account {:#034x}",
2092 event.event.account_id.0
2093 )))
2094 })?;
2095 if accepted_event.settled {
2096 entry.loss = entry.loss.checked_add(accepted_event.units).ok_or_else(|| {
2097 IngestError::Refused(StoreError(format!(
2098 "settlement loss delta overflow for account {:#034x}",
2099 event.event.account_id.0
2100 )))
2101 })?;
2102 }
2103 if accepted_event.overage {
2104 entry.overage =
2105 entry.overage.checked_add(accepted_event.units).ok_or_else(|| {
2106 IngestError::Refused(StoreError(format!(
2107 "overage delta overflow for account {:#034x}",
2108 event.event.account_id.0
2109 )))
2110 })?;
2111 }
2112 }
2113 let inserted = sqlx::query(
2114 "INSERT INTO tollgate_usage_events
2115 (request_id, account_id, lease_id, fencing_token, units, occurred_at_us, policy_revision, key_id)
2116 SELECT * FROM UNNEST($1::bytea[], $2::bytea[], $3::bytea[], $4::bigint[], $5::bigint[], $6::bigint[], $7::bytea[], $8::bytea[])",
2117 )
2118 .bind(&rid)
2119 .bind(&acct)
2120 .bind(&lease)
2121 .bind(&fence)
2122 .bind(&units)
2123 .bind(&at)
2124 .bind(&revision)
2125 .bind(&keys)
2126 .execute(&mut *tx)
2127 .await
2128 .map_err(storage)?;
2129 let expected_event_rows = u64::try_from(accepted.len())
2130 .map_err(|_| StoreError("accepted event count exceeds u64 range".into()))?;
2131 if inserted.rows_affected() != expected_event_rows {
2132 return Err(IngestError::Unavailable(StoreError(format!(
2133 "ingest inserted {} of {} accepted usage rows",
2134 inserted.rows_affected(),
2135 accepted.len()
2136 ))));
2137 }
2138
2139 let (lease_update_ids, lease_used_deltas): (Vec<Vec<u8>>, Vec<i64>) = leases
2143 .iter()
2144 .filter_map(|(lease_id, row)| {
2145 row.used_delta
2146 .map(|used_delta| (lease_id.clone(), used_delta.get()))
2147 })
2148 .unzip();
2149 if !lease_update_ids.is_empty() {
2150 let updated_leases = sqlx::query(
2151 "UPDATE tollgate_leases AS lease
2152 SET used = lease.used + delta.used
2153 FROM UNNEST($1::bytea[], $2::bigint[]) AS delta(lease_id, used)
2154 WHERE lease.lease_id = delta.lease_id",
2155 )
2156 .bind(&lease_update_ids)
2157 .bind(&lease_used_deltas)
2158 .execute(&mut *tx)
2159 .await
2160 .map_err(storage)?;
2161 let expected_lease_rows = u64::try_from(lease_update_ids.len())
2162 .map_err(|_| StoreError("ingest lease row count exceeds u64 range".into()))?;
2163 if updated_leases.rows_affected() != expected_lease_rows {
2164 return Err(IngestError::Unavailable(StoreError(format!(
2165 "ingest updated {} of {} locked lease rows",
2166 updated_leases.rows_affected(),
2167 lease_update_ids.len()
2168 ))));
2169 }
2170 }
2171
2172 let mut account_ids = Vec::with_capacity(account_deltas.len());
2173 let mut account_usage_deltas = Vec::with_capacity(account_deltas.len());
2174 let mut account_loss_deltas = Vec::with_capacity(account_deltas.len());
2175 let mut account_overage_deltas = Vec::with_capacity(account_deltas.len());
2176 for (account_id, delta) in &account_deltas {
2177 account_ids.push(account_id.clone());
2178 account_usage_deltas.push(delta.usage);
2179 account_loss_deltas.push(delta.loss);
2180 account_overage_deltas.push(delta.overage);
2181 }
2182
2183 let locked_accounts = sqlx::query(
2189 "SELECT account_id, usage_recorded, settlement_loss, overage_recorded
2190 FROM tollgate_accounts
2191 WHERE account_id = ANY($1)
2192 ORDER BY account_id FOR UPDATE",
2193 )
2194 .bind(&account_ids)
2195 .fetch_all(&mut *tx)
2196 .await
2197 .map_err(storage)?;
2198 if locked_accounts.len() != account_ids.len() {
2199 return Err(IngestError::Unavailable(StoreError(format!(
2200 "ingest locked {} of {} referenced account rows",
2201 locked_accounts.len(),
2202 account_ids.len()
2203 ))));
2204 }
2205
2206 for row in locked_accounts {
2211 let account_id: Vec<u8> = row.get(0);
2212 let usage_recorded: i64 = row.get(1);
2213 let settlement_loss: i64 = row.get(2);
2214 let overage_recorded: i64 = row.get(3);
2215 to_units(usage_recorded, "account usage_recorded")?;
2216 to_units(settlement_loss, "account settlement_loss")?;
2217 to_units(overage_recorded, "account overage_recorded")?;
2218 let delta = account_deltas.get(&account_id).ok_or_else(|| {
2219 StoreError(format!(
2220 "ingest locked unexpected account {:#034x}",
2221 id_from(&account_id)
2222 ))
2223 })?;
2224 usage_recorded.checked_add(delta.usage).ok_or_else(|| {
2225 IngestError::Refused(StoreError(format!(
2226 "usage_recorded overflow for account {:#034x}",
2227 id_from(&account_id)
2228 )))
2229 })?;
2230 overage_recorded.checked_add(delta.overage).ok_or_else(|| {
2234 IngestError::Refused(StoreError(format!(
2235 "overage_recorded overflow for account {:#034x}",
2236 id_from(&account_id)
2237 )))
2238 })?;
2239 if settlement_loss < delta.loss {
2240 return Err(IngestError::Unavailable(StoreError(format!(
2241 "settlement_loss underflow for account {:#034x}: settled straggler \
2242 usage {} exceeds recorded loss",
2243 id_from(&account_id),
2244 delta.loss
2245 ))));
2246 }
2247 }
2248
2249 let updated_accounts = sqlx::query(
2250 "UPDATE tollgate_accounts AS account
2251 SET usage_recorded = account.usage_recorded + delta.usage,
2252 settlement_loss = account.settlement_loss - delta.loss,
2253 overage_recorded = account.overage_recorded + delta.overage
2254 FROM UNNEST($1::bytea[], $2::bigint[], $3::bigint[], $4::bigint[])
2255 AS delta(account_id, usage, loss, overage)
2256 WHERE account.account_id = delta.account_id
2257 AND account.settlement_loss >= delta.loss",
2258 )
2259 .bind(&account_ids)
2260 .bind(&account_usage_deltas)
2261 .bind(&account_loss_deltas)
2262 .bind(&account_overage_deltas)
2263 .execute(&mut *tx)
2264 .await
2265 .map_err(storage)?;
2266 let expected_account_rows = u64::try_from(account_ids.len())
2267 .map_err(|_| StoreError("ingest account row count exceeds u64 range".into()))?;
2268 if updated_accounts.rows_affected() != expected_account_rows {
2269 return Err(IngestError::Unavailable(StoreError(format!(
2270 "ingest updated {} of {} locked account rows",
2271 updated_accounts.rows_affected(),
2272 account_ids.len()
2273 ))));
2274 }
2275
2276 if keys.iter().all(Option::is_none) {
2280 report.unattributed = Some(expected_event_rows);
2281 return Ok(report);
2282 }
2283
2284 let attributed: i64 = sqlx::query_scalar(
2289 "WITH matched AS MATERIALIZED (
2290 SELECT k.key_id, b.occurred_at_us
2291 FROM UNNEST($1::bytea[], $2::bytea[], $3::bigint[])
2292 AS b(key_id, account_id, occurred_at_us)
2293 JOIN LATERAL (
2294 SELECT key_id FROM tollgate_credential_keys
2295 WHERE key_id = b.key_id AND account_id = b.account_id LIMIT 1
2296 ) k ON true
2297 ), updated AS (
2298 INSERT INTO tollgate_credential_activity AS activity (key_id, last_committed_at_us)
2299 SELECT key_id, MAX(occurred_at_us) FROM matched GROUP BY key_id ORDER BY key_id
2300 ON CONFLICT (key_id) DO UPDATE
2301 SET last_committed_at_us = EXCLUDED.last_committed_at_us
2302 WHERE activity.last_committed_at_us < EXCLUDED.last_committed_at_us
2303 RETURNING key_id
2304 ) SELECT COUNT(*) FROM matched"
2305 ).bind(&keys).bind(&acct).bind(&at).fetch_one(&mut *tx).await.map_err(storage)?;
2306 report.unattributed = Some(expected_event_rows.checked_sub(
2307 u64::try_from(attributed).map_err(|_| StoreError("negative attribution count".into()))?
2308 ).ok_or_else(|| StoreError("attribution count exceeds accepted events".into()))?);
2309
2310 Ok(report)
2311 }
2312 .await;
2313 finish_transaction(tx, result).await
2317 }
2318}
2319
2320fn decode_status(stored: String) -> Result<AccountStatus, StoreError> {
2328 match stored.as_str() {
2329 s if s == AccountStatus::Active.as_str() => Ok(AccountStatus::Active),
2330 s if s == AccountStatus::Suspended.as_str() => Ok(AccountStatus::Suspended),
2331 s if s == AccountStatus::Closed.as_str() => Ok(AccountStatus::Closed),
2332 other => Err(StoreError(format!("unrecognized account status {other:?}"))),
2333 }
2334}
2335
2336fn decode_capacity_class(stored: String) -> Result<CapacityClass, StoreError> {
2345 match stored.as_str() {
2346 s if s == CapacityClass::Assured.as_str() => Ok(CapacityClass::Assured),
2347 s if s == CapacityClass::BestEffort.as_str() => Ok(CapacityClass::BestEffort),
2348 other => Err(StoreError(format!("unrecognized capacity class {other:?}"))),
2349 }
2350}
2351
2352fn decode_schedule(
2361 allowance: Option<i64>,
2362 period: Option<String>,
2363 rollover: Option<String>,
2364) -> Result<Option<BudgetSchedule>, StoreError> {
2365 let populated = [allowance.is_some(), period.is_some(), rollover.is_some()];
2366 let (Some(allowance), Some(period), Some(rollover)) = (allowance, period, rollover) else {
2367 if populated.iter().any(|present| *present) {
2368 return Err(StoreError(
2369 "stored budget schedule is partially populated".into(),
2370 ));
2371 }
2372 return Ok(None);
2373 };
2374 let period = match period.as_str() {
2375 s if s == Period::UtcCalendarMonth.as_str() => Period::UtcCalendarMonth,
2376 other => return Err(StoreError(format!("unrecognized budget period {other:?}"))),
2377 };
2378 let rollover = match rollover.as_str() {
2379 s if s == Rollover::None.as_str() => Rollover::None,
2380 other => {
2381 return Err(StoreError(format!(
2382 "unrecognized budget rollover {other:?}"
2383 )));
2384 }
2385 };
2386 Ok(Some(BudgetSchedule {
2387 allowance: to_units(allowance, "budget allowance")?,
2388 period,
2389 rollover,
2390 }))
2391}
2392
2393const BUDGET_VIEW_SQL: &str = "SELECT account_id, deposited, overage_recorded, usage_recorded,
2396 settlement_loss, expired, budget_allowance, budget_period,
2397 budget_rollover, period_start_us
2398 FROM tollgate_accounts WHERE account_id = $1";
2399
2400const GRANT_DEBIT_SQL: &str = "UPDATE tollgate_accounts
2403 SET balance = balance - $2,
2404 allowance_balance = allowance_balance - $3,
2405 next_fence = next_fence + 1
2406 WHERE account_id = $1
2407 RETURNING account_id, deposited, overage_recorded, usage_recorded,
2408 settlement_loss, expired, budget_allowance, budget_period,
2409 budget_rollover, period_start_us";
2410
2411fn budget_view(row: &sqlx::postgres::PgRow) -> Result<BudgetView, StoreError> {
2426 let deposited = to_units(row.get::<i64, _>(1), "deposited")?;
2427 let overage = to_units(row.get::<i64, _>(2), "overage_recorded")?;
2428 let usage = to_units(row.get::<i64, _>(3), "usage_recorded")?;
2429 let loss = to_units(row.get::<i64, _>(4), "settlement_loss")?;
2430 let expired = to_units(row.get::<i64, _>(5), "expired")?;
2431 let schedule = decode_schedule(row.get(6), row.get(7), row.get(8))?;
2432 let period_start = micros_ts(row.get::<i64, _>(9), "period_start_us")?;
2433
2434 let funded = deposited
2435 .checked_add(overage)
2436 .ok_or_else(|| StoreError("account funding total overflows".into()))?;
2437 let consumed = usage
2438 .checked_add(loss)
2439 .and_then(|spent| spent.checked_add(expired))
2440 .ok_or_else(|| StoreError("account consumption total overflows".into()))?;
2441 Ok(BudgetView {
2442 balance_at_publish: funded.checked_sub(consumed).ok_or_else(|| {
2443 StoreError(format!(
2444 "consumption {} exceeds funding {}",
2445 consumed.get(),
2446 funded.get()
2447 ))
2448 })?,
2449 period_end: schedule.map(|schedule| schedule.period.end_after(period_start)),
2454 })
2455}
2456
2457fn decode_publishable(
2469 principal: Principal,
2470 generation: i64,
2471 value: serde_json::Value,
2472) -> Result<PublishableSnapshot, StoreError> {
2473 let generation = generation_from(generation)?;
2474 let snapshot: StoredSnapshot =
2475 serde_json::from_value(value).map_err(|e| StoreError(format!("snapshot decode: {e}")))?;
2476 let budget = snapshot.budget;
2481 let publishable = PublishableSnapshot::try_new(Arc::new(snapshot.into_snapshot(generation)))
2482 .map_err(|error| {
2483 StoreError(format!(
2484 "invalid stored snapshot for principal {:#034x}: {error}",
2485 principal.0
2486 ))
2487 })?;
2488 Ok(match budget {
2489 Some(budget) => publishable.with_budget(Some(budget)),
2490 None => publishable,
2491 })
2492}
2493
2494fn generation_from(column: i64) -> Result<Generation, StoreError> {
2506 u64::try_from(column)
2507 .map(Generation)
2508 .map_err(|_| StoreError("stored snapshot generation is negative".into()))
2509}
2510
2511#[async_trait]
2512impl SnapshotSource for PostgresStore {
2513 async fn snapshot(&self, principal: Principal) -> Result<SnapshotResolution, StoreError> {
2514 let row = sqlx::query(
2515 "SELECT generation, snapshot, deleted FROM tollgate_snapshots WHERE principal = $1",
2516 )
2517 .bind(id_bytes(principal.0))
2518 .fetch_optional(&self.pool)
2519 .await
2520 .map_err(storage)?;
2521 match row {
2522 Some(row) if row.get::<bool, _>(2) => Ok(SnapshotResolution::Revoked {
2528 generation: generation_from(row.get::<i64, _>(0))?,
2529 }),
2530 Some(row) => Ok(SnapshotResolution::Present(decode_publishable(
2531 principal,
2532 row.get::<i64, _>(0),
2533 row.get(1),
2534 )?)),
2535 None => Ok(SnapshotResolution::Unknown),
2536 }
2537 }
2538
2539 fn subscribe(&self) -> broadcast::Receiver<SnapshotPush> {
2540 self.push.subscribe()
2541 }
2542
2543 async fn principals(&self) -> Result<Option<Vec<Principal>>, StoreError> {
2551 let rows = sqlx::query("SELECT principal FROM tollgate_snapshots ORDER BY principal")
2552 .fetch_all(&self.pool)
2553 .await
2554 .map_err(storage)?;
2555 Ok(Some(
2556 rows.iter()
2557 .map(|row| Principal(id_from(row.get::<Vec<u8>, _>(0).as_slice())))
2558 .collect(),
2559 ))
2560 }
2561}
2562
2563#[async_trait]
2564impl StoreHealth for PostgresStore {
2565 async fn ping(&self) -> Result<(), StoreError> {
2566 sqlx::query("SELECT 1")
2567 .execute(&self.pool)
2568 .await
2569 .map_err(storage)?;
2570 Ok(())
2571 }
2572}
2573
2574async fn republish_patched_snapshots(
2604 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
2605 account: AccountId,
2606 json_path: &'static str,
2607 value: &str,
2608) -> Result<(Vec<(Principal, PublishableSnapshot)>, usize), SetStatusError> {
2609 let rows = sqlx::query(
2610 "UPDATE tollgate_snapshots
2611 SET generation = generation + 1,
2612 snapshot = jsonb_set(snapshot, $3::text[], to_jsonb($2::text))
2613 WHERE account_id = $1
2614 AND deleted = FALSE
2615 AND snapshot #>> $3::text[] IS DISTINCT FROM $2::text
2616 RETURNING principal, generation, snapshot",
2617 )
2618 .bind(id_bytes(account.0))
2619 .bind(value)
2620 .bind(json_path)
2621 .fetch_all(&mut **tx)
2622 .await
2623 .map_err(storage)?;
2624
2625 let mut republished = Vec::with_capacity(rows.len());
2626 let mut unreadable = 0usize;
2627 for row in rows {
2628 let principal = Principal(id_from(row.get::<Vec<u8>, _>(0).as_slice()));
2629 match decode_publishable(
2630 principal,
2631 row.get::<i64, _>(1),
2632 row.get::<serde_json::Value, _>(2),
2633 ) {
2634 Ok(snapshot) => republished.push((principal, snapshot)),
2635 Err(error) => {
2643 unreadable += 1;
2644 tracing::warn!(
2645 %principal,
2646 %error,
2647 "restamped snapshot could not be decoded for push"
2648 );
2649 }
2650 }
2651 }
2652 republished.sort_unstable_by_key(|(principal, _)| *principal);
2655 Ok((republished, unreadable))
2656}
2657#[async_trait]
2658impl AdminStore for PostgresStore {
2659 async fn create_account(
2660 &self,
2661 config: AccountConfig,
2662 ) -> Result<tollgate_store::AdminReceipt<()>, CreateAccountError> {
2663 let result = sqlx::query(
2664 "INSERT INTO tollgate_accounts
2665 (account_id, balance, deposited, status, capacity_class, next_fence,
2666 usage_recorded, settlement_loss, overage_recorded)
2667 VALUES ($1, $2, $2, $3, $4, 1, 0, 0, 0)
2668 ON CONFLICT (account_id) DO NOTHING",
2669 )
2670 .bind(id_bytes(config.account_id.0))
2671 .bind(to_i64(config.initial_balance, "balance").map_err(CreateAccountError::Storage)?)
2672 .bind(config.status.as_str())
2673 .bind(config.capacity_class.as_str())
2674 .execute(&self.pool)
2675 .await
2676 .map_err(|e| CreateAccountError::Storage(storage(e)))?;
2677 if result.rows_affected() == 0 {
2678 return Err(CreateAccountError::AlreadyExists);
2679 }
2680 Ok(AdminReceipt::new(
2681 (),
2682 AdminState::Absent,
2683 AdminState::AccountCreated {
2684 initial_balance: config.initial_balance,
2685 status: config.status,
2686 capacity_class: config.capacity_class,
2687 },
2688 ))
2689 }
2690
2691 async fn deposit(
2692 &self,
2693 account: AccountId,
2694 units: CostUnits,
2695 ) -> Result<AdminReceipt<()>, AllocateError> {
2696 let row = sqlx::query(
2697 "UPDATE tollgate_accounts
2698 SET balance = balance + $2, deposited = deposited + $2
2699 WHERE account_id = $1
2700 RETURNING balance - $2 AS old_topup, deposited - $2 AS old_deposited,
2701 balance AS new_topup, deposited AS new_deposited",
2702 )
2703 .bind(id_bytes(account.0))
2704 .bind(i64::try_from(units.get()).map_err(|_| AllocateError::BalanceOverflow)?)
2705 .fetch_optional(&self.pool)
2706 .await
2707 .map_err(|error| {
2708 if error
2712 .as_database_error()
2713 .and_then(|db| db.code())
2714 .as_deref()
2715 == Some("22003")
2716 {
2717 AllocateError::BalanceOverflow
2718 } else {
2719 alloc_storage(error)
2720 }
2721 })?
2722 .ok_or(AllocateError::UnknownAccount)?;
2723 let state = |topup: &str, deposited: &str| -> Result<AdminState, StoreError> {
2724 Ok(AdminState::Funding {
2725 topup: to_units(row.get(topup), "audit topup")?,
2726 deposited: to_units(row.get(deposited), "audit deposited")?,
2727 })
2728 };
2729 Ok(AdminReceipt::new(
2730 (),
2731 state("old_topup", "old_deposited")?,
2732 state("new_topup", "new_deposited")?,
2733 ))
2734 }
2735
2736 async fn set_budget_schedule(
2737 &self,
2738 account: AccountId,
2739 schedule: Option<BudgetSchedule>,
2740 ) -> Result<AdminReceipt<()>, BudgetError> {
2741 let allowance = schedule
2751 .map(|s| to_i64(s.allowance, "budget allowance"))
2752 .transpose()
2753 .map_err(BudgetError::Storage)?;
2754 let mut tx = self.pool.begin().await.map_err(storage)?;
2755 let result = async {
2756 let row = sqlx::query(
2760 "SELECT budget_allowance, budget_period, budget_rollover
2761 FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE",
2762 )
2763 .bind(id_bytes(account.0))
2764 .fetch_optional(&mut *tx)
2765 .await
2766 .map_err(storage)?
2767 .ok_or(BudgetError::UnknownAccount)?;
2768 let before = decode_schedule(row.get(0), row.get(1), row.get(2))?;
2769 sqlx::query(
2770 "UPDATE tollgate_accounts
2771 SET budget_allowance = $2, budget_period = $3, budget_rollover = $4
2772 WHERE account_id = $1",
2773 )
2774 .bind(id_bytes(account.0))
2775 .bind(allowance)
2776 .bind(schedule.map(|s| s.period.as_str()))
2777 .bind(schedule.map(|s| s.rollover.as_str()))
2778 .execute(&mut *tx)
2779 .await
2780 .map_err(storage)?;
2781 Ok(AdminReceipt::new(
2782 (),
2783 AdminState::Budget { schedule: before },
2784 AdminState::Budget { schedule },
2785 ))
2786 }
2787 .await;
2788 finish_transaction(tx, result).await
2789 }
2790
2791 async fn account_view(&self, account: AccountId) -> Result<Option<AccountView>, StoreError> {
2792 let mut tx = self.pool.begin().await.map_err(storage)?;
2801 let result = async {
2802 sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY")
2803 .execute(&mut *tx)
2804 .await
2805 .map_err(storage)?;
2806 let account_row = sqlx::query(
2807 "SELECT deposited, balance, usage_recorded, settlement_loss, overage_recorded,
2808 expired, status, capacity_class, budget_allowance, budget_period,
2809 budget_rollover, period_start_us
2810 FROM tollgate_accounts WHERE account_id = $1",
2811 )
2812 .bind(id_bytes(account.0))
2813 .fetch_optional(&mut *tx)
2814 .await
2815 .map_err(storage)?;
2816 let lease_row = sqlx::query(ACTIVE_LEASE_SUM_SQL)
2817 .bind(id_bytes(account.0))
2818 .fetch_one(&mut *tx)
2819 .await
2820 .map_err(storage)?;
2821 Ok::<_, StoreError>((account_row, lease_row))
2822 }
2823 .await;
2824 let (account_row, lease_row) = finish_transaction(tx, result).await?;
2825
2826 let Some(row) = account_row else {
2827 return Ok(None);
2828 };
2829 let active_grants = to_units(lease_row.get::<i64, _>(0), "active lease grants")?;
2830 let active_used = to_units(lease_row.get::<i64, _>(1), "active lease usage")?;
2831 let recorded = to_units(row.get::<i64, _>(2), "usage_recorded")?;
2832 Ok(Some(AccountView {
2833 account_id: account,
2834 status: decode_status(row.get::<String, _>(6))?,
2835 capacity_class: decode_capacity_class(row.get::<String, _>(7))?,
2836 schedule: decode_schedule(
2837 row.get::<Option<i64>, _>(8),
2838 row.get::<Option<String>, _>(9),
2839 row.get::<Option<String>, _>(10),
2840 )?,
2841 period_start: StoredInstant {
2842 micros: row.get::<i64, _>(11),
2843 submicro_nanos: 0,
2844 }
2845 .timestamp()?,
2846 conservation: Conservation {
2847 deposited: to_units(row.get::<i64, _>(0), "deposited")?,
2848 overage_recorded: to_units(row.get::<i64, _>(4), "overage_recorded")?,
2849 balance: to_units(row.get::<i64, _>(1), "balance")?,
2850 active_lease_grants: active_grants,
2851 settled_usage: recorded.checked_sub(active_used).ok_or_else(|| {
2855 StoreError(format!(
2856 "active lease usage {} exceeds recorded usage {} for account {account}",
2857 active_used.get(),
2858 recorded.get()
2859 ))
2860 })?,
2861 settlement_loss: to_units(row.get::<i64, _>(3), "settlement_loss")?,
2862 expired: to_units(row.get::<i64, _>(5), "expired")?,
2863 },
2864 }))
2865 }
2866
2867 async fn roll_due_periods(
2868 &self,
2869 now: Timestamp,
2870 limit: NonZeroUsize,
2871 ) -> Result<RolloverBatch, StoreError> {
2872 let limit_i = i64::try_from(limit.get())
2873 .map_err(|_| StoreError(format!("rollover batch limit exceeds i64 range: {limit}")))?;
2874 let mut rolled = Vec::new();
2875 for period in Period::ALL {
2882 let boundary_us = ts_micros(period.start_of(now));
2883 let remaining = limit_i
2884 - i64::try_from(rolled.len())
2885 .map_err(|_| StoreError("rollover batch row count exceeds i64 range".into()))?;
2886 if remaining <= 0 {
2887 break;
2888 }
2889 let rows = sqlx::query(DUE_PERIODS_SQL)
2900 .bind(period.as_str())
2901 .bind(boundary_us)
2902 .bind(remaining)
2903 .fetch_all(&self.pool)
2904 .await
2905 .map_err(storage)?;
2906
2907 for row in rows {
2908 rolled.push((
2909 row.get::<i64, _>(3),
2910 RolledAccount {
2911 account_id: AccountId(id_from(&row.get::<Vec<u8>, _>(0))),
2912 deposited: to_units(row.get::<i64, _>(1), "budget allowance")?,
2913 expired: to_units(row.get::<i64, _>(2), "expiring allowance")?,
2914 },
2915 ));
2916 }
2917 }
2918 rolled
2925 .sort_unstable_by_key(|(crossed_from, account)| (*crossed_from, account.account_id.0));
2926 RolloverBatch::try_new(
2927 rolled.into_iter().map(|(_, account)| account).collect(),
2928 limit,
2929 )
2930 }
2931
2932 async fn set_account_status(
2933 &self,
2934 account: AccountId,
2935 status: AccountStatus,
2936 ) -> Result<tollgate_store::AdminReceipt<StatusChange>, SetStatusError> {
2937 let mut tx = self.pool.begin().await.map_err(storage)?;
2938 let result = async {
2939 let row = sqlx::query(
2944 "SELECT status FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE",
2945 )
2946 .bind(id_bytes(account.0))
2947 .fetch_optional(&mut *tx)
2948 .await
2949 .map_err(storage)?
2950 .ok_or(SetStatusError::UnknownAccount)?;
2951
2952 let before = decode_status(row.get::<String, _>(0))?;
2953 if before == AccountStatus::Closed && status != AccountStatus::Closed {
2954 return Err(SetStatusError::AccountClosed);
2957 }
2958
2959 sqlx::query("UPDATE tollgate_accounts SET status = $2 WHERE account_id = $1")
2960 .bind(id_bytes(account.0))
2961 .bind(status.as_str())
2962 .execute(&mut *tx)
2963 .await
2964 .map_err(storage)?;
2965
2966 let (republished, unreadable) =
2967 republish_patched_snapshots(&mut tx, account, "{status}", status.as_str()).await?;
2968 Ok((before, republished, unreadable))
2969 }
2970 .await;
2971
2972 let (before, republished, unreadable) = finish_transaction(tx, result).await?;
2973 if pushes_exceed_capacity(republished.len()) {
2975 tracing::warn!(
2976 %account,
2977 principals = republished.len(),
2978 capacity = PUSH_CHANNEL_CAPACITY,
2979 "status change emitted more pushes than the channel holds; subscribers will resync"
2980 );
2981 }
2982 let count = republished.len();
2983 for (principal, snapshot) in republished {
2984 self.push_to_subscribers(SnapshotPush {
2985 principal,
2986 resolution: SnapshotResolution::Present(snapshot),
2987 });
2988 }
2989 Ok(AdminReceipt::new(
2990 StatusChange {
2991 republished: count + unreadable,
2994 unreadable,
2995 },
2996 AdminState::Status { status: before },
2997 AdminState::Status { status },
2998 ))
2999 }
3000
3001 async fn set_capacity_class(
3002 &self,
3003 account: AccountId,
3004 class: CapacityClass,
3005 ) -> Result<tollgate_store::AdminReceipt<StatusChange>, SetStatusError> {
3006 let mut tx = self.pool.begin().await.map_err(storage)?;
3007 let result = async {
3008 let row = sqlx::query(
3016 "SELECT status, capacity_class FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE",
3017 )
3018 .bind(id_bytes(account.0))
3019 .fetch_optional(&mut *tx)
3020 .await
3021 .map_err(storage)?
3022 .ok_or(SetStatusError::UnknownAccount)?;
3023
3024 let before = decode_capacity_class(row.get::<String, _>(1))?;
3025 if decode_status(row.get::<String, _>(0))? == AccountStatus::Closed {
3030 return Err(SetStatusError::AccountClosed);
3031 }
3032
3033 sqlx::query("UPDATE tollgate_accounts SET capacity_class = $2 WHERE account_id = $1")
3034 .bind(id_bytes(account.0))
3035 .bind(class.as_str())
3036 .execute(&mut *tx)
3037 .await
3038 .map_err(storage)?;
3039
3040 let (republished, unreadable) =
3041 republish_patched_snapshots(&mut tx, account, "{capacity_class}", class.as_str()).await?;
3042 Ok((before, republished, unreadable))
3043 }
3044 .await;
3045
3046 let (before, republished, unreadable) = finish_transaction(tx, result).await?;
3047 if pushes_exceed_capacity(republished.len()) {
3048 tracing::warn!(
3049 %account,
3050 principals = republished.len(),
3051 capacity = PUSH_CHANNEL_CAPACITY,
3052 "capacity class change emitted more pushes than the channel holds; \
3053 subscribers will resync"
3054 );
3055 }
3056 let count = republished.len();
3057 for (principal, snapshot) in republished {
3058 self.push_to_subscribers(SnapshotPush {
3059 principal,
3060 resolution: SnapshotResolution::Present(snapshot),
3061 });
3062 }
3063 Ok(AdminReceipt::new(
3064 StatusChange {
3065 republished: count + unreadable,
3068 unreadable,
3069 },
3070 AdminState::CapacityClass {
3071 capacity_class: before,
3072 },
3073 AdminState::CapacityClass {
3074 capacity_class: class,
3075 },
3076 ))
3077 }
3078
3079 async fn publish_snapshot(
3080 &self,
3081 principal: Principal,
3082 snapshot: PublishableSnapshot,
3083 ) -> Result<tollgate_store::AdminReceipt<()>, PublishSnapshotError> {
3084 let generation = i64::try_from(snapshot.generation.0).map_err(|_| {
3085 StoreError("snapshot generation exceeds PostgreSQL BIGINT range".into())
3086 })?;
3087
3088 let mut tx = self.pool.begin().await.map_err(storage)?;
3089 let result = publish_in_tx(&mut tx, principal, generation, snapshot).await;
3090
3091 let (written, published, before, after) = finish_transaction(tx, result).await?;
3092 if written {
3093 self.push_to_subscribers(SnapshotPush {
3097 principal,
3098 resolution: SnapshotResolution::Present(published),
3099 });
3100 }
3101 Ok(AdminReceipt::new((), before, after))
3102 }
3103
3104 async fn remove_snapshot(&self, principal: Principal) -> Result<AdminReceipt<()>, StoreError> {
3105 let mut tx = self.pool.begin().await.map_err(storage)?;
3106 let result = remove_in_tx(&mut tx, principal).await;
3107 let receipt = finish_transaction(tx, result).await?;
3108 self.announce_removal(principal, &receipt);
3109 Ok(receipt)
3110 }
3111}
3112
3113async fn lock_account_key(
3137 tx: &mut Transaction<'_, Postgres>,
3138 account: AccountId,
3139 key: KeyId,
3140) -> Result<(Principal, bool), KeySnapshotError> {
3141 let row = sqlx::query(
3142 "SELECT principal, revoked_at_us IS NOT NULL
3143 FROM tollgate_credential_keys WHERE key_id = $1 AND account_id = $2 FOR SHARE",
3144 )
3145 .bind(id_bytes(key.0))
3146 .bind(id_bytes(account.0))
3147 .fetch_optional(&mut **tx)
3148 .await
3149 .map_err(storage)?
3150 .ok_or(KeySnapshotError::UnknownCredential)?;
3151 let principal: Vec<u8> = row.get(0);
3152 let principal: [u8; 16] = principal
3153 .try_into()
3154 .map_err(|_| StoreError("credential principal is not 16 bytes".into()))?;
3155 Ok((Principal(u128::from_be_bytes(principal)), row.get(1)))
3156}
3157
3158async fn publish_in_tx(
3164 tx: &mut Transaction<'_, Postgres>,
3165 principal: Principal,
3166 generation: i64,
3167 snapshot: PublishableSnapshot,
3168) -> Result<(bool, PublishableSnapshot, AdminState, AdminState), PublishSnapshotError> {
3169 if let Some(key_id) = snapshot.key_id {
3170 let matches: bool = sqlx::query_scalar(
3171 "SELECT EXISTS(SELECT 1 FROM tollgate_credential_keys
3172 WHERE key_id = $1 AND principal = $2 AND account_id = $3)",
3173 )
3174 .bind(id_bytes(key_id.0))
3175 .bind(id_bytes(principal.0))
3176 .bind(id_bytes(snapshot.account_id.0))
3177 .fetch_one(&mut **tx)
3178 .await
3179 .map_err(storage)?;
3180 if !matches {
3181 return Err(PublishSnapshotError::CredentialMismatch { key_id });
3182 }
3183 }
3184 let ledger = sqlx::query(
3196 "SELECT status, deposited, overage_recorded, usage_recorded, settlement_loss,
3197 expired, budget_allowance, budget_period, budget_rollover, period_start_us,
3198 capacity_class
3199 FROM tollgate_accounts WHERE account_id = $1 FOR SHARE",
3200 )
3201 .bind(id_bytes(snapshot.account_id.0))
3202 .fetch_optional(&mut **tx)
3203 .await
3204 .map_err(storage)?;
3205
3206 let view = match &ledger {
3214 Some(row) => {
3215 let ledger = decode_status(row.get::<String, _>(0))?;
3216 if ledger != snapshot.status {
3217 return Err(PublishSnapshotError::StatusMismatch {
3218 ledger,
3219 submitted: snapshot.status,
3220 });
3221 }
3222 let ledger_class = decode_capacity_class(row.get::<String, _>(10))?;
3226 if ledger_class != snapshot.capacity_class {
3227 return Err(PublishSnapshotError::CapacityClassMismatch {
3228 ledger: ledger_class,
3229 submitted: snapshot.capacity_class,
3230 });
3231 }
3232 Some(budget_view(row)?)
3233 }
3234 None => None,
3235 };
3236 let published = snapshot.with_budget(view);
3237 let value = serde_json::to_value(StoredSnapshotRef::from(published.as_snapshot()))
3238 .map_err(|e| StoreError(format!("snapshot encode: {e}")))?;
3239
3240 let (written, before, after) = write_snapshot_audited(tx, principal, generation, value).await?;
3241 Ok((written, published, before, after))
3242}
3243
3244async fn remove_in_tx(
3247 tx: &mut Transaction<'_, Postgres>,
3248 principal: Principal,
3249) -> Result<AdminReceipt<()>, StoreError> {
3250 let before = snapshot_audit_row(tx, principal).await?;
3251 let after = match before {
3252 AdminState::Snapshot {
3253 generation,
3254 revoked: false,
3255 } => {
3256 sqlx::query("UPDATE tollgate_snapshots SET deleted = TRUE WHERE principal = $1")
3257 .bind(id_bytes(principal.0))
3258 .execute(&mut **tx)
3259 .await
3260 .map_err(storage)?;
3261 AdminState::Snapshot {
3262 generation,
3263 revoked: true,
3264 }
3265 }
3266 state => state,
3267 };
3268 Ok(AdminReceipt::new((), before, after))
3269}
3270
3271async fn snapshot_audit_row(
3274 tx: &mut Transaction<'_, Postgres>,
3275 principal: Principal,
3276) -> Result<AdminState, StoreError> {
3277 let row = sqlx::query(
3278 "SELECT generation, deleted FROM tollgate_snapshots WHERE principal = $1 FOR UPDATE",
3279 )
3280 .bind(id_bytes(principal.0))
3281 .fetch_optional(&mut **tx)
3282 .await
3283 .map_err(storage)?;
3284 row.map(|row| {
3285 Ok(AdminState::Snapshot {
3286 generation: generation_from(row.get(0))?,
3287 revoked: row.get(1),
3288 })
3289 })
3290 .unwrap_or(Ok(AdminState::Absent))
3291}
3292
3293async fn write_snapshot_audited(
3294 tx: &mut Transaction<'_, Postgres>,
3295 principal: Principal,
3296 generation: i64,
3297 value: serde_json::Value,
3298) -> Result<(bool, AdminState, AdminState), StoreError> {
3299 let mut before = snapshot_audit_row(tx, principal).await?;
3300 let after = AdminState::Snapshot {
3301 generation: generation_from(generation)?,
3302 revoked: false,
3303 };
3304 if before == AdminState::Absent {
3305 let inserted = sqlx::query(
3306 "INSERT INTO tollgate_snapshots (principal, generation, snapshot, deleted)
3307 VALUES ($1, $2, $3, FALSE) ON CONFLICT (principal) DO NOTHING",
3308 )
3309 .bind(id_bytes(principal.0))
3310 .bind(generation)
3311 .bind(&value)
3312 .execute(&mut **tx)
3313 .await
3314 .map_err(storage)?;
3315 if inserted.rows_affected() == 1 {
3316 return Ok((true, before, after));
3317 }
3318 before = snapshot_audit_row(tx, principal).await?;
3321 }
3322 let AdminState::Snapshot {
3323 generation: previous,
3324 ..
3325 } = before
3326 else {
3327 return Err(StoreError("snapshot disappeared during publication".into()));
3328 };
3329 if previous >= generation_from(generation)? {
3330 return Ok((false, before, before));
3331 }
3332 sqlx::query("UPDATE tollgate_snapshots SET generation = $2, snapshot = $3, deleted = FALSE WHERE principal = $1")
3333 .bind(id_bytes(principal.0)).bind(generation).bind(value)
3334 .execute(&mut **tx).await.map_err(storage)?;
3335 Ok((true, before, after))
3336}
3337
3338#[cfg(test)]
3339mod tests {
3340 use super::*;
3341
3342 #[test]
3343 fn stored_fences_use_the_exact_positive_bigint_domain() {
3344 for invalid in [i64::MIN, -1, 0] {
3345 assert!(stored_fence(invalid).is_err());
3346 }
3347 for valid in [1, 2, i64::MAX] {
3348 assert_eq!(stored_fence(valid).unwrap(), FencingToken(valid as u64));
3349 }
3350 }
3351
3352 #[tokio::test]
3355 async fn ping_surfaces_a_closed_pool() {
3356 let pool = PgPoolOptions::new()
3357 .connect_lazy("postgres://localhost/tollgate")
3358 .unwrap();
3359 pool.close().await;
3360 let (push, _) = broadcast::channel(1);
3361 let store = PostgresStore {
3362 pool,
3363 policy: GrantPolicy::default(),
3364 push,
3365 };
3366
3367 assert!(store.ping().await.is_err());
3368 assert!(matches!(
3369 AdminStore::deposit(&store, AccountId(1), CostUnits(1)).await,
3370 Err(AllocateError::Storage(_))
3371 ));
3372 }
3373}