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>, KeyError> {
329 validate_key_page_limit(limit)?;
330 let sql = if after.is_some() {
335 "SELECT page.* FROM tollgate_accounts AS account
336 LEFT JOIN LATERAL (
337 SELECT key_id, not_after_floor_us, not_after_submicro_ns,
338 not_after_is_lower_bound, revoked_at_us
339 FROM tollgate_credential_keys
340 WHERE account_id = account.account_id AND key_id > $3
341 ORDER BY key_id LIMIT $2
342 ) AS page ON TRUE
343 WHERE account.account_id = $1 ORDER BY page.key_id"
344 } else {
345 "SELECT page.* FROM tollgate_accounts AS account
346 LEFT JOIN LATERAL (
347 SELECT key_id, not_after_floor_us, not_after_submicro_ns,
348 not_after_is_lower_bound, revoked_at_us
349 FROM tollgate_credential_keys
350 WHERE account_id = account.account_id
351 ORDER BY key_id LIMIT $2
352 ) AS page ON TRUE
353 WHERE account.account_id = $1 ORDER BY page.key_id"
354 };
355 let mut query = sqlx::query(sql)
356 .bind(id_bytes(account.0))
357 .bind(i64::try_from(limit.get()).unwrap_or(i64::MAX));
358 if let Some(cursor) = after {
359 query = query.bind(id_bytes(cursor.0));
360 }
361 let rows = query.fetch_all(&self.pool).await.map_err(storage)?;
362 if rows.is_empty() {
363 return Err(KeyError::UnknownAccount);
364 }
365 rows.into_iter()
366 .filter(|row| row.get::<Option<&[u8]>, _>(0).is_some())
367 .map(|row| summary_from_row(row).map_err(KeyError::Storage))
368 .collect()
369 }
370
371 async fn insert_key_within(
372 &self,
373 record: KeyRecord,
374 max_active: NonZeroUsize,
375 now: Timestamp,
376 ) -> Result<(), KeyError> {
377 self.insert_key_within_audited(record, max_active, now)
378 .await
379 .map(|receipt| receipt.outcome)
380 }
381
382 async fn insert_key_within_audited(
383 &self,
384 record: KeyRecord,
385 max_active: NonZeroUsize,
386 now: Timestamp,
387 ) -> Result<AdminReceipt<()>, KeyError> {
388 self.insert_credential(record, Some((max_active, now)))
389 .await
390 }
391}
392
393impl PostgresStore {
394 async fn insert_credential(
399 &self,
400 record: KeyRecord,
401 bound: Option<(NonZeroUsize, Timestamp)>,
402 ) -> Result<AdminReceipt<()>, KeyError> {
403 let mut tx = self
404 .pool
405 .begin()
406 .await
407 .map_err(|e| KeyError::Storage(storage(e)))?;
408 let account = sqlx::query(ACCOUNT_LOCK_SQL)
416 .bind(id_bytes(record.account_id.0))
417 .fetch_optional(&mut *tx)
418 .await
419 .map_err(|e| KeyError::Storage(storage(e)))?;
420 if account.is_none() {
421 return Err(KeyError::UnknownAccount);
422 }
423
424 if let Some((max_active, now)) = bound {
425 let existing: Option<i32> = sqlx::query_scalar(
432 "SELECT 1 FROM tollgate_credential_keys WHERE key_id = $1 OR principal = $2",
433 )
434 .bind(id_bytes(record.key_id.0))
435 .bind(id_bytes(record.principal.0))
436 .fetch_optional(&mut *tx)
437 .await
438 .map_err(|e| KeyError::Storage(storage(e)))?;
439 if existing.is_some() {
440 return Err(KeyError::AlreadyExists);
441 }
442
443 let cutoff = StoredInstant::from(now);
444 let live: i64 = sqlx::query_scalar(LIVE_KEY_COUNT_SQL)
445 .bind(id_bytes(record.account_id.0))
446 .bind(cutoff.micros)
447 .bind(cutoff.submicro_nanos)
448 .fetch_one(&mut *tx)
449 .await
450 .map_err(|e| KeyError::Storage(storage(e)))?;
451 if u128::from(live.max(0).unsigned_abs())
452 >= u128::try_from(max_active.get()).unwrap_or(u128::MAX)
453 {
454 return Err(KeyError::ActiveKeyLimit { limit: max_active });
455 }
456 }
457
458 let expiry = record.not_after.map(StoredInstant::from);
459 let result = sqlx::query(
460 "INSERT INTO tollgate_credential_keys
461 (key_id, account_id, principal, digest, not_after_floor_us,
462 not_after_submicro_ns, not_after_is_lower_bound, revoked_at_us)
463 VALUES ($1, $2, $3, $4, $5, $6, FALSE, NULL)
464 ON CONFLICT (key_id) DO NOTHING",
465 )
466 .bind(id_bytes(record.key_id.0))
467 .bind(id_bytes(record.account_id.0))
468 .bind(id_bytes(record.principal.0))
469 .bind(record.digest.to_vec())
470 .bind(expiry.map(|expiry| expiry.micros))
471 .bind(expiry.map(|expiry| expiry.submicro_nanos))
472 .execute(&mut *tx)
473 .await;
474 let outcome = match result {
475 Ok(done) if done.rows_affected() == 0 => Err(KeyError::AlreadyExists),
476 Ok(_) => Ok(()),
477 Err(sqlx::Error::Database(e)) if e.is_unique_violation() => {
478 Err(KeyError::AlreadyExists)
479 }
480 Err(e) => Err(KeyError::Storage(storage(e))),
490 };
491 outcome?;
492 tx.commit()
493 .await
494 .map_err(|e| KeyError::Storage(storage(e)))?;
495 Ok(AdminReceipt::new(
496 (),
497 AdminState::Absent,
498 AdminState::Credential {
499 account_id: record.account_id,
500 key_id: record.key_id,
501 revoked: false,
502 },
503 ))
504 }
505}
506
507fn summary_from_row(row: sqlx::postgres::PgRow) -> Result<KeySummary, StoreError> {
508 let bytes: Vec<u8> = row.get(0);
509 let fixed: [u8; 16] = bytes
510 .as_slice()
511 .try_into()
512 .map_err(|_| StoreError("credential identifier is not 16 bytes".into()))?;
513 let not_after = match (row.get::<Option<i64>, _>(1), row.get::<Option<i16>, _>(2)) {
514 (None, None) if !row.get::<bool, _>(3) => None,
515 (Some(micros), Some(submicro_nanos)) => Some(
516 StoredInstant {
517 micros,
518 submicro_nanos,
519 }
520 .timestamp()?,
521 ),
522 _ => return Err(StoreError("incomplete stored credential expiry".into())),
523 };
524 Ok(KeySummary {
525 key_id: KeyId(u128::from_be_bytes(fixed)),
526 not_after,
527 revoked_at: row
528 .get::<Option<i64>, _>(4)
529 .map(|micros| {
530 StoredInstant {
531 micros,
532 submicro_nanos: 0,
533 }
534 .timestamp()
535 })
536 .transpose()?,
537 })
538}
539
540fn credential_from_row(row: sqlx::postgres::PgRow) -> Result<KeyRecord, StoreError> {
541 let id = |index| -> Result<u128, StoreError> {
542 let bytes: Vec<u8> = row.get(index);
543 let fixed: [u8; 16] = bytes
544 .as_slice()
545 .try_into()
546 .map_err(|_| StoreError("credential identifier is not 16 bytes".into()))?;
547 Ok(u128::from_be_bytes(fixed))
548 };
549 let key_id = KeyId(id(0)?);
550 let not_after = match (row.get::<Option<i64>, _>(4), row.get::<Option<i16>, _>(5)) {
551 (None, None) if !row.get::<bool, _>(6) => None,
552 (Some(micros), Some(submicro_nanos)) => Some(
553 StoredInstant {
554 micros,
555 submicro_nanos,
556 }
557 .timestamp()?,
558 ),
559 _ => return Err(StoreError("incomplete stored credential expiry".into())),
560 };
561 Ok(KeyRecord {
562 key_id,
563 account_id: AccountId(id(1)?),
564 principal: Principal(id(2)?),
565 digest: digest_from(row.get::<Vec<u8>, _>(3).as_slice(), key_id)?,
566 not_after,
567 })
568}
569
570#[async_trait]
571impl tollgate_store::KeySource for PostgresStore {
572 async fn active_keys_page(
573 &self,
574 now: Timestamp,
575 after: Option<KeyId>,
576 limit: NonZeroUsize,
577 ) -> Result<tollgate_store::KeyPage, StoreError> {
578 validate_key_page_limit(limit)?;
579 let cutoff = StoredInstant::from(now);
580 let mut tx = self.pool.begin().await.map_err(storage)?;
583 sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY")
584 .execute(&mut *tx)
585 .await
586 .map_err(storage)?;
587 let revision: i64 =
588 sqlx::query_scalar("SELECT revision FROM tollgate_credential_revision WHERE singleton")
589 .fetch_one(&mut *tx)
590 .await
591 .map_err(storage)?;
592 let revision = u64::try_from(revision)
593 .map_err(|_| StoreError("credential revision is negative".into()))?;
594 let sql = if after.is_some() {
596 "SELECT key_id, account_id, principal, digest, not_after_floor_us, not_after_submicro_ns, not_after_is_lower_bound
597 FROM tollgate_credential_keys WHERE revoked_at_us IS NULL
598 AND (not_after_floor_us IS NULL OR not_after_submicro_ns IS NULL
599 OR (not_after_floor_us, not_after_submicro_ns) > ($1, $2)) AND key_id > $3
600 ORDER BY key_id LIMIT $4"
601 } else {
602 "SELECT key_id, account_id, principal, digest, not_after_floor_us, not_after_submicro_ns, not_after_is_lower_bound
603 FROM tollgate_credential_keys WHERE revoked_at_us IS NULL
604 AND (not_after_floor_us IS NULL OR not_after_submicro_ns IS NULL
605 OR (not_after_floor_us, not_after_submicro_ns) > ($1, $2)) AND key_id >= $3
606 ORDER BY key_id LIMIT $4"
607 };
608 let rows = sqlx::query(sql)
609 .bind(cutoff.micros)
610 .bind(cutoff.submicro_nanos)
611 .bind(id_bytes(after.unwrap_or(KeyId(0)).0))
612 .bind((limit.get() + 1) as i64)
613 .fetch_all(&mut *tx)
614 .await
615 .map_err(storage)?;
616 tx.commit().await.map_err(storage)?;
617 let mut records = rows
620 .into_iter()
621 .map(credential_from_row)
622 .map(|row| row.map(tollgate_store::CredentialRecord::from))
623 .collect::<Result<Vec<_>, _>>()?;
624 let next_after = if records.len() > limit.get() {
625 records.pop();
626 records.last().map(|record| record.key_id)
627 } else {
628 None
629 };
630 tollgate_store::KeyPage::try_new(revision, now, after, limit, records, next_after)
631 }
632}
633
634#[cfg(test)]
635mod stored_id_tests {
636 use super::StoredId;
637
638 #[test]
639 fn malformed_storage_id_explains_both_accepted_representations() {
640 let error = serde_json::from_value::<StoredId>(serde_json::Value::Bool(true)).unwrap_err();
641 assert!(
642 error
643 .to_string()
644 .contains("a legacy u64 number or a canonical 128-bit identifier string"),
645 "unexpected diagnostic: {error}"
646 );
647 }
648}
649
650#[derive(Serialize)]
666struct StoredSnapshotRef<'a> {
667 account_id: StoredId,
668 key_id: Option<StoredId>,
669 status: &'a AccountStatus,
670 capacity_class: &'a CapacityClass,
673 enforcement_mode: &'a EnforcementMode,
674 valid_until: &'a Timestamp,
675 permissions: &'a PermissionBits,
676 limits: &'a ResolvedLimits,
677 cost_table: &'a Arc<CostTable>,
678 budget: Option<&'a BudgetView>,
682 policy_revision: &'a PolicyRevision,
687}
688
689impl<'a> From<&'a AccountSnapshot> for StoredSnapshotRef<'a> {
690 fn from(snapshot: &'a AccountSnapshot) -> Self {
691 StoredSnapshotRef {
692 account_id: StoredId(snapshot.account_id.0),
693 key_id: snapshot.key_id.map(|id| StoredId(id.0)),
694 status: &snapshot.status,
695 capacity_class: &snapshot.capacity_class,
696 enforcement_mode: &snapshot.enforcement_mode,
697 valid_until: &snapshot.valid_until,
698 permissions: &snapshot.permissions,
699 limits: &snapshot.limits,
700 cost_table: &snapshot.cost_table,
701 budget: snapshot.budget.as_ref(),
702 policy_revision: &snapshot.policy_revision,
703 }
704 }
705}
706
707#[derive(Deserialize)]
711struct StoredSnapshot {
712 account_id: StoredId,
713 key_id: Option<StoredId>,
714 status: AccountStatus,
715 #[serde(default)]
721 capacity_class: CapacityClass,
722 #[serde(default)]
733 enforcement_mode: EnforcementMode,
734 valid_until: Timestamp,
735 permissions: PermissionBits,
736 limits: ResolvedLimits,
737 cost_table: Arc<CostTable>,
738 #[serde(default)]
744 budget: Option<BudgetView>,
745 #[serde(default)]
751 policy_revision: PolicyRevision,
752}
753
754impl StoredSnapshot {
755 fn into_snapshot(self, generation: Generation) -> AccountSnapshot {
761 let builder = AccountSnapshot::builder(
762 AccountId(self.account_id.0),
763 generation,
764 self.status,
765 self.valid_until,
766 self.permissions,
767 self.limits,
768 self.cost_table,
769 )
770 .enforcement_mode(self.enforcement_mode)
771 .capacity_class(self.capacity_class)
772 .policy_revision(self.policy_revision);
773 match self.key_id {
774 Some(key_id) => builder.key_id(tollgate_core::KeyId(key_id.0)).build(),
775 None => builder.build(),
776 }
777 }
778}
779
780const ACTIVE_LEASE_SUM_SQL: &str =
798 "SELECT COALESCE(SUM(granted), 0)::BIGINT, COALESCE(SUM(used), 0)::BIGINT
799 FROM tollgate_leases WHERE account_id = $1 AND state = 0";
800
801pub(crate) const ACCOUNT_LOCK_SQL: &str =
820 "SELECT 1 FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE";
821
822pub(crate) const LIVE_KEY_COUNT_SQL: &str = "SELECT count(*) FROM tollgate_credential_keys
825 WHERE account_id = $1
826 AND revoked_at_us IS NULL
827 AND (not_after_floor_us IS NULL OR not_after_submicro_ns IS NULL
828 OR (not_after_floor_us, not_after_submicro_ns) > ($2, $3))";
829
830const RECLAIM_DUE_LEASES_SQL: &str = "SELECT lease_id, account_id, granted, used
831 FROM tollgate_leases
832 WHERE state = 0 AND (expires_at_floor_us, expires_at_submicro_ns) <= ($1, $2)
833 ORDER BY expires_at_floor_us, expires_at_submicro_ns
834 LIMIT $3 FOR UPDATE SKIP LOCKED";
835
836const DUE_PERIODS_SQL: &str = "WITH due AS (
865 SELECT account_id, allowance_balance AS prior, budget_allowance AS allowance,
866 period_start_us AS crossed_from
867 FROM tollgate_accounts
868 WHERE budget_allowance IS NOT NULL
869 AND budget_period = $1 AND period_start_us < $2
870 ORDER BY period_start_us LIMIT $3 FOR UPDATE SKIP LOCKED
871 )
872 UPDATE tollgate_accounts AS account SET
873 deposited = account.deposited + due.allowance,
874 expired = account.expired + due.prior,
875 balance = account.balance - due.prior + due.allowance,
876 allowance_balance = due.allowance,
877 period_start_us = $2
878 FROM due
879 WHERE account.account_id = due.account_id
880 RETURNING due.account_id, due.allowance, due.prior, due.crossed_from";
881
882fn id_bytes(id: u128) -> Vec<u8> {
883 id.to_be_bytes().to_vec()
884}
885
886fn id_from(bytes: &[u8]) -> u128 {
887 let mut buf = [0u8; 16];
888 buf.copy_from_slice(bytes);
889 u128::from_be_bytes(buf)
890}
891
892fn to_i64(units: CostUnits, what: &str) -> Result<i64, StoreError> {
893 i64::try_from(units.get()).map_err(|_| StoreError(format!("{what} exceeds i64 range")))
894}
895
896fn stored_fence(value: i64) -> Result<FencingToken, StoreError> {
899 u64::try_from(value)
900 .ok()
901 .filter(|value| *value > 0)
902 .map(FencingToken)
903 .ok_or_else(|| StoreError(format!("stored fencing token is not positive: {value}")))
904}
905
906fn to_units(value: i64, what: &str) -> Result<CostUnits, StoreError> {
910 u64::try_from(value)
911 .map(CostUnits)
912 .map_err(|_| StoreError(format!("{what} is negative in storage: {value}")))
913}
914
915fn ts_micros(ts: Timestamp) -> i64 {
916 ts.as_microsecond()
917}
918
919fn micros_ts(value: i64, what: &str) -> Result<Timestamp, StoreError> {
922 tollgate_store::clock::timestamp_from_micros(value).map_err(|e| {
923 StoreError(format!(
924 "{what} is not a representable instant: {value} ({e})"
925 ))
926 })
927}
928
929fn storage(e: sqlx::Error) -> StoreError {
930 StoreError(format!("postgres: {e}"))
931}
932
933fn alloc_storage(e: sqlx::Error) -> AllocateError {
934 AllocateError::Storage(storage(e))
935}
936
937async fn finish_transaction<T, E>(
950 tx: Transaction<'_, Postgres>,
951 result: Result<T, E>,
952) -> Result<T, E>
953where
954 E: From<StoreError> + std::fmt::Display,
955{
956 match result {
957 Ok(value) => {
958 tx.commit().await.map_err(|e| E::from(storage(e)))?;
959 Ok(value)
960 }
961 Err(error) => {
962 if let Err(rollback_error) = tx.rollback().await {
963 return Err(E::from(StoreError(format!(
964 "operation failed ({error}); transaction rollback failed ({})",
965 storage(rollback_error)
966 ))));
967 }
968 Err(error)
969 }
970 }
971}
972
973pub struct PostgresStore {
1010 pool: PgPool,
1011 policy: GrantPolicy,
1012 push: broadcast::Sender<SnapshotPush>,
1013}
1014
1015#[derive(Debug, Clone, Copy)]
1020pub struct PoolConfig {
1021 pub max_connections: u32,
1024 pub acquire_timeout: std::time::Duration,
1027}
1028
1029impl Default for PoolConfig {
1030 fn default() -> Self {
1031 PoolConfig {
1032 max_connections: 16,
1033 acquire_timeout: std::time::Duration::from_secs(5),
1034 }
1035 }
1036}
1037
1038impl PoolConfig {
1039 pub fn validate(&self) -> Result<(), StoreError> {
1045 if self.max_connections == 0 {
1046 return Err(StoreError("max_connections must be positive".into()));
1047 }
1048 if self.acquire_timeout.is_zero() {
1049 return Err(StoreError("acquire_timeout must be positive".into()));
1050 }
1051 Ok(())
1052 }
1053}
1054
1055impl PostgresStore {
1056 pub async fn connect(url: &str, policy: GrantPolicy) -> Result<Arc<Self>, StoreError> {
1058 Self::connect_with(url, policy, PoolConfig::default()).await
1059 }
1060
1061 pub async fn connect_with(
1070 url: &str,
1071 policy: GrantPolicy,
1072 pool_config: PoolConfig,
1073 ) -> Result<Arc<Self>, StoreError> {
1074 policy
1075 .validate()
1076 .map_err(|e| StoreError(format!("invalid grant policy: {e}")))?;
1077 pool_config.validate()?;
1078 let pool = PgPoolOptions::new()
1079 .max_connections(pool_config.max_connections)
1080 .acquire_timeout(pool_config.acquire_timeout)
1081 .connect(url)
1082 .await
1083 .map_err(storage)?;
1084 sqlx::migrate!("./migrations")
1085 .run(&pool)
1086 .await
1087 .map_err(|e| StoreError(format!("migrate: {e}")))?;
1088 let (push, _) = broadcast::channel(PUSH_CHANNEL_CAPACITY);
1089 Ok(Arc::new(PostgresStore { pool, policy, push }))
1090 }
1091
1092 fn push_to_subscribers(&self, push: SnapshotPush) {
1098 let principal = push.principal;
1099 let subscribers = self.push.send(push).unwrap_or(0);
1100 tracing::debug!(
1101 %principal,
1102 subscribers,
1103 "snapshot pushed to subscribers"
1104 );
1105 }
1106
1107 fn announce_removal(&self, principal: Principal, receipt: &AdminReceipt<()>) {
1109 if receipt.before != receipt.after
1110 && let AdminState::Snapshot { generation, .. } = receipt.after
1111 {
1112 self.push_to_subscribers(SnapshotPush {
1113 principal,
1114 resolution: SnapshotResolution::Revoked { generation },
1115 });
1116 }
1117 }
1118
1119 pub async fn balance(&self, account: AccountId) -> Result<CostUnits, StoreError> {
1129 let row = sqlx::query("SELECT balance FROM tollgate_accounts WHERE account_id = $1")
1130 .bind(id_bytes(account.0))
1131 .fetch_optional(&self.pool)
1132 .await
1133 .map_err(storage)?;
1134 row.map(|r| to_units(r.get::<i64, _>(0), "account balance"))
1135 .transpose()
1136 .map(|units| units.unwrap_or(CostUnits::ZERO))
1137 }
1138
1139 pub async fn usage_recorded(&self, account: AccountId) -> Result<CostUnits, StoreError> {
1146 let row = sqlx::query("SELECT usage_recorded FROM tollgate_accounts WHERE account_id = $1")
1147 .bind(id_bytes(account.0))
1148 .fetch_optional(&self.pool)
1149 .await
1150 .map_err(storage)?;
1151 row.map(|r| to_units(r.get::<i64, _>(0), "account usage_recorded"))
1152 .transpose()
1153 .map(|units| units.unwrap_or(CostUnits::ZERO))
1154 }
1155
1156 pub async fn conservation(
1171 &self,
1172 account: AccountId,
1173 ) -> Result<Option<Conservation>, StoreError> {
1174 let mut tx = self.pool.begin().await.map_err(storage)?;
1188 let result = async {
1189 sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY")
1190 .execute(&mut *tx)
1191 .await
1192 .map_err(storage)?;
1193 let account_row = sqlx::query(
1194 "SELECT deposited, balance, usage_recorded, settlement_loss, overage_recorded,
1195 expired
1196 FROM tollgate_accounts WHERE account_id = $1",
1197 )
1198 .bind(id_bytes(account.0))
1199 .fetch_optional(&mut *tx)
1200 .await
1201 .map_err(storage)?;
1202 let lease_row = sqlx::query(ACTIVE_LEASE_SUM_SQL)
1203 .bind(id_bytes(account.0))
1204 .fetch_one(&mut *tx)
1205 .await
1206 .map_err(storage)?;
1207 Ok::<_, StoreError>((account_row, lease_row))
1208 }
1209 .await;
1210 let (account_row, lease_row) = finish_transaction(tx, result).await?;
1211
1212 let Some(row) = account_row else {
1213 return Ok(None);
1214 };
1215 let active_grants = to_units(lease_row.get::<i64, _>(0), "active lease grants")?;
1216 let active_used = to_units(lease_row.get::<i64, _>(1), "active lease usage")?;
1217 let recorded = to_units(row.get::<i64, _>(2), "usage_recorded")?;
1218 Ok(Some(Conservation {
1219 deposited: to_units(row.get::<i64, _>(0), "deposited")?,
1220 overage_recorded: to_units(row.get::<i64, _>(4), "overage_recorded")?,
1221 balance: to_units(row.get::<i64, _>(1), "balance")?,
1222 active_lease_grants: active_grants,
1223 settled_usage: recorded.checked_sub(active_used).ok_or_else(|| {
1231 StoreError(format!(
1232 "active lease usage {} exceeds recorded usage {} for account {account}",
1233 active_used.get(),
1234 recorded.get()
1235 ))
1236 })?,
1237 settlement_loss: to_units(row.get::<i64, _>(3), "settlement_loss")?,
1238 expired: to_units(row.get::<i64, _>(5), "expired")?,
1239 }))
1240 }
1241}
1242
1243struct LockedLeaseRow {
1247 account_id: Vec<u8>,
1248 fencing_token: i64,
1249 granted: i64,
1250 used: i64,
1251 credited: i64,
1252 expires_at: StoredInstant,
1253 state: i16,
1254 from_allowance: i64,
1259 period_start_us: i64,
1260}
1261
1262struct ReleasedCredit {
1265 account: AccountId,
1266 restored: CostUnits,
1267 preserves_funding: bool,
1268}
1269
1270async fn lock_lease(
1272 tx: &mut Transaction<'_, Postgres>,
1273 lease_id: LeaseId,
1274) -> Result<Option<LockedLeaseRow>, sqlx::Error> {
1275 let row = sqlx::query(
1276 "SELECT account_id, fencing_token, granted, used, credited, expires_at_floor_us, state,
1277 from_allowance, period_start_us, expires_at_submicro_ns
1278 FROM tollgate_leases WHERE lease_id = $1 FOR UPDATE",
1279 )
1280 .bind(id_bytes(lease_id.0))
1281 .fetch_optional(&mut **tx)
1282 .await?;
1283 Ok(row.map(|row| LockedLeaseRow {
1284 account_id: row.get(0),
1285 fencing_token: row.get(1),
1286 granted: row.get(2),
1287 used: row.get(3),
1288 credited: row.get(4),
1289 expires_at: StoredInstant {
1290 micros: row.get(5),
1291 submicro_nanos: row.get(9),
1292 },
1293 state: row.get(6),
1294 from_allowance: row.get(7),
1295 period_start_us: row.get(8),
1296 }))
1297}
1298
1299struct Exchange {
1302 floor: CostUnits,
1306 needed: CostUnits,
1309 preserves_funding: bool,
1312}
1313
1314impl Exchange {
1315 const ACQUIRE: Exchange = Exchange {
1316 floor: CostUnits::ZERO,
1317 needed: CostUnits::ZERO,
1318 preserves_funding: true,
1319 };
1320}
1321
1322impl PostgresStore {
1323 async fn acquire_in_tx(
1325 &self,
1326 tx: &mut Transaction<'_, Postgres>,
1327 account: AccountId,
1328 requested: CostUnits,
1329 expires_at: Timestamp,
1330 exchange: Exchange,
1331 ) -> Result<Allocation, AllocateError> {
1332 let Exchange {
1333 floor,
1334 needed,
1335 preserves_funding: settlement_preserves_funding,
1336 } = exchange;
1337 let row = sqlx::query(
1338 "SELECT balance, status, next_fence, allowance_balance, period_start_us
1339 FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE",
1340 )
1341 .bind(id_bytes(account.0))
1342 .fetch_optional(&mut **tx)
1343 .await
1344 .map_err(alloc_storage)?
1345 .ok_or(AllocateError::UnknownAccount)?;
1346
1347 if decode_status(row.get::<String, _>(1)).map_err(AllocateError::Storage)?
1351 != AccountStatus::Active
1352 {
1353 return Err(AllocateError::AccountInactive);
1354 }
1355 let balance =
1356 to_units(row.get::<i64, _>(0), "account balance").map_err(AllocateError::Storage)?;
1357 let granted = match self
1363 .policy
1364 .consolidation_grant(requested, balance, floor, needed)
1365 {
1366 Some(granted) => granted,
1367 None => {
1368 if settlement_preserves_funding && !requested.is_zero() {
1371 let budget = sqlx::query(BUDGET_VIEW_SQL)
1372 .bind(id_bytes(account.0))
1373 .fetch_one(&mut **tx)
1374 .await
1375 .map_err(alloc_storage)?;
1376 let evidence = budget_view(&budget)
1377 .map_err(AllocateError::Storage)?
1378 .shortfall();
1379 return Err(match evidence.exhaustion() {
1380 Some(exhausted) => AllocateError::BalanceExhausted(exhausted),
1381 None => AllocateError::BalanceInsufficient(evidence),
1382 });
1383 }
1384 return Err(AllocateError::InsufficientBalance);
1385 }
1386 };
1387 let granted_i = to_i64(granted, "grant").map_err(AllocateError::Storage)?;
1392 let allowance_balance = row.get::<i64, _>(3);
1393 let from_allowance = granted_i.min(allowance_balance);
1394 let period_start_us = row.get::<i64, _>(4);
1395 let fence = row.get::<i64, _>(2);
1396 let fence_token = stored_fence(fence).map_err(AllocateError::Storage)?;
1397
1398 let lease_id = LeaseId(uuid::Uuid::new_v4().as_u128());
1399
1400 let budget = sqlx::query(GRANT_DEBIT_SQL)
1404 .bind(id_bytes(account.0))
1405 .bind(granted_i)
1406 .bind(from_allowance)
1407 .fetch_one(&mut **tx)
1408 .await
1409 .map_err(alloc_storage)?;
1410 let funding = budget_view(&budget)
1411 .map_err(AllocateError::Storage)?
1412 .shortfall();
1413 sqlx::query(
1414 "INSERT INTO tollgate_leases
1415 (lease_id, account_id, fencing_token, granted, used, credited, expires_at_floor_us,
1416 state, from_allowance, period_start_us, expires_at_submicro_ns, expiry_is_upper_bound)
1417 VALUES ($1, $2, $3, $4, 0, 0, $5, 0, $6, $7, $8, FALSE)",
1418 )
1419 .bind(id_bytes(lease_id.0))
1420 .bind(id_bytes(account.0))
1421 .bind(fence)
1422 .bind(granted_i)
1423 .bind(StoredInstant::from(expires_at).micros)
1424 .bind(from_allowance)
1425 .bind(period_start_us)
1426 .bind(StoredInstant::from(expires_at).submicro_nanos)
1427 .execute(&mut **tx)
1428 .await
1429 .map_err(alloc_storage)?;
1430
1431 Ok(Allocation {
1432 grant: LeaseGrant {
1433 lease_id,
1434 account_id: account,
1435 fencing_token: fence_token,
1436 units: granted,
1437 expires_at,
1438 },
1439 funding: Some(funding),
1440 })
1441 }
1442
1443 async fn release_in_tx(
1446 &self,
1447 tx: &mut Transaction<'_, Postgres>,
1448 lease_id: LeaseId,
1449 fencing_token: FencingToken,
1450 unspent: CostUnits,
1451 now: Timestamp,
1452 ) -> Result<ReleasedCredit, AllocateError> {
1453 let LockedLeaseRow {
1454 account_id,
1455 fencing_token: fence,
1456 granted,
1457 used,
1458 credited: _credited,
1459 expires_at,
1460 state,
1461 from_allowance,
1462 period_start_us,
1463 } = lock_lease(tx, lease_id)
1464 .await
1465 .map_err(alloc_storage)?
1466 .ok_or(AllocateError::UnknownLease)?;
1467
1468 if stored_fence(fence).map_err(AllocateError::Storage)? != fencing_token {
1469 return Err(AllocateError::Fenced);
1470 }
1471 let expires_at = expires_at.timestamp().map_err(AllocateError::Storage)?;
1474 if state != STATE_ACTIVE
1475 || self
1476 .policy
1477 .reclaim_cutoff(now)
1478 .is_some_and(|cutoff| expires_at <= cutoff)
1479 {
1480 return Err(AllocateError::LeaseNotActive);
1481 }
1482 let account = AccountId(id_from(&account_id));
1486 let unspent_i = to_i64(unspent, "unspent").map_err(AllocateError::Storage)?;
1487 let loss = granted
1488 .checked_sub(
1489 used.checked_add(unspent_i)
1490 .ok_or(AllocateError::InvalidRelease)?,
1491 )
1492 .filter(|l| *l >= 0)
1493 .ok_or(AllocateError::InvalidRelease)?;
1494
1495 sqlx::query("UPDATE tollgate_leases SET state = $2, credited = $3 WHERE lease_id = $1")
1496 .bind(id_bytes(lease_id.0))
1497 .bind(STATE_RELEASED)
1498 .bind(unspent_i)
1499 .execute(&mut **tx)
1500 .await
1501 .map_err(alloc_storage)?;
1502 let from_topup = granted
1508 .checked_sub(from_allowance)
1509 .filter(|t| *t >= 0)
1510 .ok_or_else(|| {
1511 AllocateError::Storage(StoreError(format!(
1512 "lease allowance funding {from_allowance} exceeds its grant {granted}"
1513 )))
1514 })?;
1515 let to_topup = from_topup.min(unspent_i);
1516 let to_allowance = unspent_i - to_topup;
1517
1518 let restored: i64 = sqlx::query_scalar(
1525 "UPDATE tollgate_accounts SET
1526 balance = balance + $2
1527 + CASE WHEN period_start_us > $5 THEN 0 ELSE $3 END,
1528 allowance_balance = allowance_balance
1529 + CASE WHEN period_start_us > $5 THEN 0 ELSE $3 END,
1530 expired = expired + CASE WHEN period_start_us > $5 THEN $3 ELSE 0 END,
1531 settlement_loss = settlement_loss + $4
1532 WHERE account_id = $1
1533 RETURNING $2 + CASE WHEN period_start_us > $5 THEN 0 ELSE $3 END",
1534 )
1535 .bind(account_id)
1536 .bind(to_topup)
1537 .bind(to_allowance)
1538 .bind(loss)
1539 .bind(period_start_us)
1540 .fetch_one(&mut **tx)
1541 .await
1542 .map_err(alloc_storage)?;
1543 Ok(ReleasedCredit {
1544 account,
1545 preserves_funding: loss == 0 && restored == unspent_i,
1546 restored: to_units(restored, "restored release credit")
1547 .map_err(AllocateError::Storage)?,
1548 })
1549 }
1550
1551 fn grant_expiry(
1555 &self,
1556 ttl: SignedDuration,
1557 now: Timestamp,
1558 ) -> Result<Timestamp, AllocateError> {
1559 if ttl <= SignedDuration::ZERO {
1560 return Err(AllocateError::InvalidTtl);
1561 }
1562 now.checked_add(ttl.min(self.policy.max_ttl))
1563 .map_err(|e| AllocateError::Storage(StoreError(format!("ttl overflow: {e}"))))
1564 }
1565}
1566
1567#[async_trait]
1568impl LeaseAllocator for PostgresStore {
1569 async fn acquire(
1570 &self,
1571 account: AccountId,
1572 requested: CostUnits,
1573 ttl: SignedDuration,
1574 now: Timestamp,
1575 ) -> Result<Allocation, AllocateError> {
1576 let expires_at = self.grant_expiry(ttl, now)?;
1577 let mut tx = self.pool.begin().await.map_err(alloc_storage)?;
1578 let result = self
1579 .acquire_in_tx(&mut tx, account, requested, expires_at, Exchange::ACQUIRE)
1580 .await;
1581 finish_transaction(tx, result).await
1582 }
1583
1584 async fn release(
1585 &self,
1586 lease_id: LeaseId,
1587 fencing_token: FencingToken,
1588 unspent: CostUnits,
1589 now: Timestamp,
1590 ) -> Result<(), AllocateError> {
1591 let mut tx = self.pool.begin().await.map_err(alloc_storage)?;
1592 let result = self
1593 .release_in_tx(&mut tx, lease_id, fencing_token, unspent, now)
1594 .await
1595 .map(|_| ());
1596 finish_transaction(tx, result).await
1597 }
1598
1599 async fn consolidate(
1600 &self,
1601 lease_id: LeaseId,
1602 fencing_token: FencingToken,
1603 unspent: CostUnits,
1604 requested: CostUnits,
1605 needed: CostUnits,
1606 ttl: SignedDuration,
1607 now: Timestamp,
1608 ) -> Result<Allocation, AllocateError> {
1609 let expires_at = self.grant_expiry(ttl, now)?;
1610 let mut tx = self.pool.begin().await.map_err(alloc_storage)?;
1611 let result = async {
1620 let released = self
1621 .release_in_tx(&mut tx, lease_id, fencing_token, unspent, now)
1622 .await?;
1623 self.acquire_in_tx(
1624 &mut tx,
1625 released.account,
1626 requested,
1627 expires_at,
1628 Exchange {
1629 floor: released.restored,
1630 needed,
1631 preserves_funding: released.preserves_funding,
1632 },
1633 )
1634 .await
1635 }
1636 .await;
1637 finish_transaction(tx, result).await
1638 }
1639
1640 async fn reclaim_expired_batch(
1641 &self,
1642 now: Timestamp,
1643 limit: NonZeroUsize,
1644 ) -> Result<ReclaimBatch, StoreError> {
1645 let limit_i = i64::try_from(limit.get())
1646 .map_err(|_| StoreError(format!("reclaim batch limit exceeds i64 range: {limit}")))?;
1647 let Some(cutoff) = self.policy.reclaim_cutoff(now) else {
1648 return ReclaimBatch::try_new(Vec::new(), limit);
1649 };
1650 let cutoff = StoredInstant::from(cutoff);
1651 let mut tx = self.pool.begin().await.map_err(storage)?;
1652 let result = async {
1653 let rows = sqlx::query(RECLAIM_DUE_LEASES_SQL)
1659 .bind(cutoff.micros)
1660 .bind(cutoff.submicro_nanos)
1661 .bind(limit_i)
1662 .fetch_all(&mut *tx)
1663 .await
1664 .map_err(storage)?;
1665
1666 let mut reclaimed = Vec::with_capacity(rows.len());
1671 let mut lease_ids = Vec::with_capacity(rows.len());
1672 let mut forfeits: std::collections::BTreeMap<Vec<u8>, i64> =
1673 std::collections::BTreeMap::new();
1674 for row in rows {
1675 let lease_bytes: Vec<u8> = row.get(0);
1676 let account_bytes: Vec<u8> = row.get(1);
1677 let granted = row.get::<i64, _>(2);
1678 let used = row.get::<i64, _>(3);
1679 let forfeited = granted.checked_sub(used).ok_or_else(|| {
1683 StoreError(format!(
1684 "reclaim remainder overflow: granted {granted}, used {used}"
1685 ))
1686 })?;
1687 let forfeited_units = to_units(forfeited, "reclaim remainder")?;
1688 let total = forfeits.entry(account_bytes.clone()).or_default();
1689 *total = total
1690 .checked_add(forfeited)
1691 .ok_or_else(|| StoreError("reclaim loss sum overflow".into()))?;
1692 lease_ids.push(lease_bytes.clone());
1693 reclaimed.push(ReclaimedLease {
1694 lease_id: LeaseId(id_from(&lease_bytes)),
1695 account_id: AccountId(id_from(&account_bytes)),
1696 forfeited: forfeited_units,
1697 });
1698 }
1699
1700 let batch = ReclaimBatch::try_new(reclaimed, limit)?;
1701 if batch.is_empty() {
1702 return Ok(batch);
1703 }
1704
1705 let (account_ids, account_forfeits): (Vec<_>, Vec<_>) = forfeits.into_iter().unzip();
1706 let expected_lease_rows = u64::try_from(lease_ids.len())
1707 .map_err(|_| StoreError("reclaim lease row count exceeds u64 range".into()))?;
1708 let expected_account_rows = u64::try_from(account_ids.len())
1709 .map_err(|_| StoreError("reclaim account row count exceeds u64 range".into()))?;
1710
1711 let locked_accounts = sqlx::query(
1717 "SELECT account_id FROM tollgate_accounts
1718 WHERE account_id = ANY($1) ORDER BY account_id FOR UPDATE",
1719 )
1720 .bind(&account_ids)
1721 .fetch_all(&mut *tx)
1722 .await
1723 .map_err(storage)?;
1724 if locked_accounts.len() != account_ids.len() {
1725 return Err(StoreError(format!(
1726 "reclaim locked {} of {} referenced account rows",
1727 locked_accounts.len(),
1728 account_ids.len()
1729 )));
1730 }
1731
1732 let updated_leases = sqlx::query(
1733 "UPDATE tollgate_leases
1734 SET state = $2, credited = 0
1735 WHERE lease_id = ANY($1) AND state = $3",
1736 )
1737 .bind(&lease_ids)
1738 .bind(STATE_EXPIRED)
1739 .bind(STATE_ACTIVE)
1740 .execute(&mut *tx)
1741 .await
1742 .map_err(storage)?;
1743 if updated_leases.rows_affected() != expected_lease_rows {
1744 return Err(StoreError(format!(
1745 "reclaim updated {} of {} locked lease rows",
1746 updated_leases.rows_affected(),
1747 lease_ids.len()
1748 )));
1749 }
1750
1751 let updated_accounts = sqlx::query(
1752 "UPDATE tollgate_accounts AS account
1753 SET settlement_loss = account.settlement_loss + delta.forfeited
1754 FROM UNNEST($1::bytea[], $2::bigint[]) AS delta(account_id, forfeited)
1755 WHERE account.account_id = delta.account_id",
1756 )
1757 .bind(&account_ids)
1758 .bind(&account_forfeits)
1759 .execute(&mut *tx)
1760 .await
1761 .map_err(storage)?;
1762 if updated_accounts.rows_affected() != expected_account_rows {
1763 return Err(StoreError(format!(
1764 "reclaim updated {} of {} locked account rows",
1765 updated_accounts.rows_affected(),
1766 account_ids.len()
1767 )));
1768 }
1769
1770 Ok(batch)
1771 }
1772 .await;
1773 finish_transaction(tx, result).await
1774 }
1775}
1776
1777#[async_trait]
1778impl UsageSink for PostgresStore {
1779 async fn ingest(
1780 &self,
1781 events: &[UsageEvent],
1782 _now: Timestamp,
1783 ) -> Result<IngestReport, IngestError> {
1784 let mut report = IngestReport {
1792 unattributed: Some(0),
1793 ..IngestReport::default()
1794 };
1795 if events.is_empty() {
1796 return Ok(report);
1797 }
1798
1799 struct PreparedEvent<'a> {
1805 event: &'a UsageEvent,
1806 request_id: Vec<u8>,
1807 account_id: Vec<u8>,
1808 key_id: Option<Vec<u8>>,
1809 lease_id: Option<Vec<u8>>,
1813 occurred_at_us: i64,
1814 }
1815 let prepared: Vec<PreparedEvent<'_>> = events
1816 .iter()
1817 .map(|event| {
1818 Ok(PreparedEvent {
1819 event,
1820 request_id: id_bytes(event.request_id.0),
1821 account_id: id_bytes(event.account_id.0),
1822 key_id: event.key_id.map(|id| id_bytes(id.0)),
1823 lease_id: event.source.lease_id().map(|id| id_bytes(id.0)),
1824 occurred_at_us: ts_micros(event.occurred_at),
1825 })
1826 })
1827 .collect::<Result<_, StoreError>>()?;
1828
1829 let mut tx = self.pool.begin().await.map_err(storage)?;
1830 let result: Result<_, IngestError> = async {
1831 let lease_ids: Vec<Vec<u8>> = prepared
1840 .iter()
1841 .filter_map(|event| event.lease_id.clone())
1842 .collect::<std::collections::BTreeSet<_>>()
1843 .into_iter()
1844 .collect();
1845 struct LeaseRow {
1846 account_id: Vec<u8>,
1847 fence: i64,
1848 granted: i64,
1849 used: i64,
1850 used_delta: Option<NonZeroI64>,
1851 credited: i64,
1852 settled: bool,
1853 }
1854 let rows = sqlx::query(
1855 "SELECT lease_id, account_id, fencing_token, granted, used, credited, state
1856 FROM tollgate_leases
1857 WHERE lease_id = ANY($1)
1858 ORDER BY account_id, lease_id FOR UPDATE",
1859 )
1860 .bind(&lease_ids)
1861 .fetch_all(&mut *tx)
1862 .await
1863 .map_err(storage)?;
1864 let mut leases: std::collections::BTreeMap<Vec<u8>, LeaseRow> =
1865 std::collections::BTreeMap::new();
1866 for row in rows {
1867 let lease_id: Vec<u8> = row.get(0);
1868 let fence: i64 = row.get(2);
1869 let granted: i64 = row.get(3);
1870 let used: i64 = row.get(4);
1871 let credited: i64 = row.get(5);
1872 leases.insert(
1873 lease_id,
1874 LeaseRow {
1875 account_id: row.get(1),
1876 fence,
1877 granted,
1878 used,
1879 used_delta: None,
1880 credited,
1881 settled: row.get::<i16, _>(6) != STATE_ACTIVE,
1882 },
1883 );
1884 }
1885
1886 let request_ids: Vec<Vec<u8>> = prepared
1888 .iter()
1889 .map(|event| event.request_id.clone())
1890 .collect();
1891 let mut seen: std::collections::HashSet<Vec<u8>> = sqlx::query(
1892 "SELECT request_id FROM tollgate_usage_events WHERE request_id = ANY($1)",
1893 )
1894 .bind(&request_ids)
1895 .fetch_all(&mut *tx)
1896 .await
1897 .map_err(storage)?
1898 .into_iter()
1899 .map(|row| row.get::<Vec<u8>, _>(0))
1900 .collect();
1901
1902 let overage_account_ids: Vec<Vec<u8>> = prepared
1918 .iter()
1919 .filter(|event| event.lease_id.is_none())
1920 .map(|event| event.account_id.clone())
1921 .collect::<std::collections::BTreeSet<_>>()
1922 .into_iter()
1923 .collect();
1924 let known_overage_accounts: std::collections::HashSet<Vec<u8>> =
1925 if overage_account_ids.is_empty() {
1926 std::collections::HashSet::new()
1927 } else {
1928 sqlx::query(
1929 "SELECT account_id FROM tollgate_accounts WHERE account_id = ANY($1)",
1930 )
1931 .bind(&overage_account_ids)
1932 .fetch_all(&mut *tx)
1933 .await
1934 .map_err(storage)?
1935 .into_iter()
1936 .map(|row| row.get::<Vec<u8>, _>(0))
1937 .collect()
1938 };
1939
1940 struct Accepted {
1944 event_index: usize,
1945 settled: bool,
1946 fence: Option<i64>,
1951 overage: bool,
1954 units: i64,
1957 }
1958 let mut accepted: Vec<Accepted> = Vec::with_capacity(prepared.len());
1959 for (event_index, event) in prepared.iter().enumerate() {
1960 if seen.contains(event.request_id.as_slice()) {
1961 report.duplicate += 1;
1962 continue;
1963 }
1964 let Some(lease_key) = event.lease_id.as_deref() else {
1965 if !known_overage_accounts.contains(event.account_id.as_slice()) {
1974 report.rejected += 1;
1975 continue;
1976 }
1977 let Ok(units) = i64::try_from(event.event.units.get()) else {
1978 report.rejected += 1;
1981 continue;
1982 };
1983 seen.insert(event.request_id.clone());
1984 accepted.push(Accepted {
1985 event_index,
1986 settled: false,
1990 fence: None,
1991 overage: true,
1992 units,
1993 });
1994 report.accepted += 1;
1995 continue;
1996 };
1997 let Some(lease) = leases.get_mut(lease_key) else {
1998 report.rejected += 1;
1999 continue;
2000 };
2001 if Some(stored_fence(lease.fence)?) != event.event.source.fencing_token()
2002 || lease.account_id.as_slice() != event.account_id.as_slice()
2003 {
2004 report.rejected += 1;
2005 continue;
2006 }
2007 to_units(lease.granted, "lease granted")?;
2008 to_units(lease.used, "lease used")?;
2009 to_units(lease.credited, "lease credited")?;
2010 let Ok(units) = i64::try_from(event.event.units.get()) else {
2011 report.rejected += 1;
2012 continue;
2013 };
2014 let committed = lease
2015 .used
2016 .checked_add(lease.used_delta.map_or(0, NonZeroI64::get))
2017 .and_then(|used| used.checked_add(lease.credited))
2018 .ok_or_else(|| {
2019 StoreError(format!(
2020 "lease accounting overflow for {:#034x}",
2021 id_from(lease_key)
2022 ))
2023 })?;
2024 let remaining = lease.granted.checked_sub(committed).ok_or_else(|| {
2025 StoreError(format!(
2026 "lease accounting exceeds grant for {:#034x}: granted {}, committed {committed}",
2027 id_from(lease_key), lease.granted
2028 ))
2029 })?;
2030 if units > remaining {
2031 report.rejected += 1;
2032 continue;
2033 }
2034 let used_delta = lease
2035 .used_delta
2036 .map_or(0, NonZeroI64::get)
2037 .checked_add(units)
2038 .ok_or_else(|| {
2039 StoreError(format!(
2040 "lease usage delta overflow for {:#034x}",
2041 id_from(lease_key)
2042 ))
2043 })?;
2044 lease.used_delta = NonZeroI64::new(used_delta);
2045 seen.insert(event.request_id.clone());
2046 accepted.push(Accepted {
2047 event_index,
2048 settled: lease.settled,
2049 fence: Some(lease.fence),
2050 overage: false,
2051 units,
2052 });
2053 report.accepted += 1;
2054 }
2055
2056 if accepted.is_empty() {
2057 return Ok(report);
2058 }
2059
2060 #[derive(Default)]
2061 struct AccountDelta {
2062 usage: i64,
2063 loss: i64,
2064 overage: i64,
2068 }
2069
2070 let (mut rid, mut acct, mut lease, mut fence, mut units, mut at, mut revision, mut keys) = (
2072 Vec::with_capacity(accepted.len()),
2073 Vec::with_capacity(accepted.len()),
2074 Vec::with_capacity(accepted.len()),
2075 Vec::with_capacity(accepted.len()),
2076 Vec::with_capacity(accepted.len()),
2077 Vec::with_capacity(accepted.len()),
2078 Vec::with_capacity(accepted.len()),
2079 Vec::with_capacity(accepted.len()),
2080 );
2081 let mut account_deltas: std::collections::BTreeMap<Vec<u8>, AccountDelta> =
2082 std::collections::BTreeMap::new();
2083 for accepted_event in &accepted {
2084 let event = &prepared[accepted_event.event_index];
2085 rid.push(event.request_id.clone());
2086 acct.push(event.account_id.clone());
2087 keys.push(event.key_id.clone());
2088 lease.push(event.lease_id.clone());
2089 fence.push(accepted_event.fence);
2090 units.push(accepted_event.units);
2091 at.push(event.occurred_at_us);
2092 revision.push(event.event.policy_revision.as_bytes().to_vec());
2097 if accepted_event.overage {
2098 debug_assert!(
2099 event.lease_id.is_none() && accepted_event.fence.is_none(),
2100 "an overage row must carry neither half of a capability"
2101 );
2102 }
2103
2104 let entry = account_deltas.entry(event.account_id.clone()).or_default();
2105 entry.usage = entry.usage.checked_add(accepted_event.units).ok_or_else(|| {
2106 IngestError::Refused(StoreError(format!(
2107 "usage delta overflow for account {:#034x}",
2108 event.event.account_id.0
2109 )))
2110 })?;
2111 if accepted_event.settled {
2112 entry.loss = entry.loss.checked_add(accepted_event.units).ok_or_else(|| {
2113 IngestError::Refused(StoreError(format!(
2114 "settlement loss delta overflow for account {:#034x}",
2115 event.event.account_id.0
2116 )))
2117 })?;
2118 }
2119 if accepted_event.overage {
2120 entry.overage =
2121 entry.overage.checked_add(accepted_event.units).ok_or_else(|| {
2122 IngestError::Refused(StoreError(format!(
2123 "overage delta overflow for account {:#034x}",
2124 event.event.account_id.0
2125 )))
2126 })?;
2127 }
2128 }
2129 let inserted = sqlx::query(
2130 "INSERT INTO tollgate_usage_events
2131 (request_id, account_id, lease_id, fencing_token, units, occurred_at_us, policy_revision, key_id)
2132 SELECT * FROM UNNEST($1::bytea[], $2::bytea[], $3::bytea[], $4::bigint[], $5::bigint[], $6::bigint[], $7::bytea[], $8::bytea[])",
2133 )
2134 .bind(&rid)
2135 .bind(&acct)
2136 .bind(&lease)
2137 .bind(&fence)
2138 .bind(&units)
2139 .bind(&at)
2140 .bind(&revision)
2141 .bind(&keys)
2142 .execute(&mut *tx)
2143 .await
2144 .map_err(storage)?;
2145 let expected_event_rows = u64::try_from(accepted.len())
2146 .map_err(|_| StoreError("accepted event count exceeds u64 range".into()))?;
2147 if inserted.rows_affected() != expected_event_rows {
2148 return Err(IngestError::Unavailable(StoreError(format!(
2149 "ingest inserted {} of {} accepted usage rows",
2150 inserted.rows_affected(),
2151 accepted.len()
2152 ))));
2153 }
2154
2155 let (lease_update_ids, lease_used_deltas): (Vec<Vec<u8>>, Vec<i64>) = leases
2159 .iter()
2160 .filter_map(|(lease_id, row)| {
2161 row.used_delta
2162 .map(|used_delta| (lease_id.clone(), used_delta.get()))
2163 })
2164 .unzip();
2165 if !lease_update_ids.is_empty() {
2166 let updated_leases = sqlx::query(
2167 "UPDATE tollgate_leases AS lease
2168 SET used = lease.used + delta.used
2169 FROM UNNEST($1::bytea[], $2::bigint[]) AS delta(lease_id, used)
2170 WHERE lease.lease_id = delta.lease_id",
2171 )
2172 .bind(&lease_update_ids)
2173 .bind(&lease_used_deltas)
2174 .execute(&mut *tx)
2175 .await
2176 .map_err(storage)?;
2177 let expected_lease_rows = u64::try_from(lease_update_ids.len())
2178 .map_err(|_| StoreError("ingest lease row count exceeds u64 range".into()))?;
2179 if updated_leases.rows_affected() != expected_lease_rows {
2180 return Err(IngestError::Unavailable(StoreError(format!(
2181 "ingest updated {} of {} locked lease rows",
2182 updated_leases.rows_affected(),
2183 lease_update_ids.len()
2184 ))));
2185 }
2186 }
2187
2188 let mut account_ids = Vec::with_capacity(account_deltas.len());
2189 let mut account_usage_deltas = Vec::with_capacity(account_deltas.len());
2190 let mut account_loss_deltas = Vec::with_capacity(account_deltas.len());
2191 let mut account_overage_deltas = Vec::with_capacity(account_deltas.len());
2192 for (account_id, delta) in &account_deltas {
2193 account_ids.push(account_id.clone());
2194 account_usage_deltas.push(delta.usage);
2195 account_loss_deltas.push(delta.loss);
2196 account_overage_deltas.push(delta.overage);
2197 }
2198
2199 let locked_accounts = sqlx::query(
2205 "SELECT account_id, usage_recorded, settlement_loss, overage_recorded
2206 FROM tollgate_accounts
2207 WHERE account_id = ANY($1)
2208 ORDER BY account_id FOR UPDATE",
2209 )
2210 .bind(&account_ids)
2211 .fetch_all(&mut *tx)
2212 .await
2213 .map_err(storage)?;
2214 if locked_accounts.len() != account_ids.len() {
2215 return Err(IngestError::Unavailable(StoreError(format!(
2216 "ingest locked {} of {} referenced account rows",
2217 locked_accounts.len(),
2218 account_ids.len()
2219 ))));
2220 }
2221
2222 for row in locked_accounts {
2227 let account_id: Vec<u8> = row.get(0);
2228 let usage_recorded: i64 = row.get(1);
2229 let settlement_loss: i64 = row.get(2);
2230 let overage_recorded: i64 = row.get(3);
2231 to_units(usage_recorded, "account usage_recorded")?;
2232 to_units(settlement_loss, "account settlement_loss")?;
2233 to_units(overage_recorded, "account overage_recorded")?;
2234 let delta = account_deltas.get(&account_id).ok_or_else(|| {
2235 StoreError(format!(
2236 "ingest locked unexpected account {:#034x}",
2237 id_from(&account_id)
2238 ))
2239 })?;
2240 usage_recorded.checked_add(delta.usage).ok_or_else(|| {
2241 IngestError::Refused(StoreError(format!(
2242 "usage_recorded overflow for account {:#034x}",
2243 id_from(&account_id)
2244 )))
2245 })?;
2246 overage_recorded.checked_add(delta.overage).ok_or_else(|| {
2250 IngestError::Refused(StoreError(format!(
2251 "overage_recorded overflow for account {:#034x}",
2252 id_from(&account_id)
2253 )))
2254 })?;
2255 if settlement_loss < delta.loss {
2256 return Err(IngestError::Unavailable(StoreError(format!(
2257 "settlement_loss underflow for account {:#034x}: settled straggler \
2258 usage {} exceeds recorded loss",
2259 id_from(&account_id),
2260 delta.loss
2261 ))));
2262 }
2263 }
2264
2265 let updated_accounts = sqlx::query(
2266 "UPDATE tollgate_accounts AS account
2267 SET usage_recorded = account.usage_recorded + delta.usage,
2268 settlement_loss = account.settlement_loss - delta.loss,
2269 overage_recorded = account.overage_recorded + delta.overage
2270 FROM UNNEST($1::bytea[], $2::bigint[], $3::bigint[], $4::bigint[])
2271 AS delta(account_id, usage, loss, overage)
2272 WHERE account.account_id = delta.account_id
2273 AND account.settlement_loss >= delta.loss",
2274 )
2275 .bind(&account_ids)
2276 .bind(&account_usage_deltas)
2277 .bind(&account_loss_deltas)
2278 .bind(&account_overage_deltas)
2279 .execute(&mut *tx)
2280 .await
2281 .map_err(storage)?;
2282 let expected_account_rows = u64::try_from(account_ids.len())
2283 .map_err(|_| StoreError("ingest account row count exceeds u64 range".into()))?;
2284 if updated_accounts.rows_affected() != expected_account_rows {
2285 return Err(IngestError::Unavailable(StoreError(format!(
2286 "ingest updated {} of {} locked account rows",
2287 updated_accounts.rows_affected(),
2288 account_ids.len()
2289 ))));
2290 }
2291
2292 if keys.iter().all(Option::is_none) {
2296 report.unattributed = Some(expected_event_rows);
2297 return Ok(report);
2298 }
2299
2300 let attributed: i64 = sqlx::query_scalar(
2305 "WITH matched AS MATERIALIZED (
2306 SELECT k.key_id, b.occurred_at_us
2307 FROM UNNEST($1::bytea[], $2::bytea[], $3::bigint[])
2308 AS b(key_id, account_id, occurred_at_us)
2309 JOIN LATERAL (
2310 SELECT key_id FROM tollgate_credential_keys
2311 WHERE key_id = b.key_id AND account_id = b.account_id LIMIT 1
2312 ) k ON true
2313 ), updated AS (
2314 INSERT INTO tollgate_credential_activity AS activity (key_id, last_committed_at_us)
2315 SELECT key_id, MAX(occurred_at_us) FROM matched GROUP BY key_id ORDER BY key_id
2316 ON CONFLICT (key_id) DO UPDATE
2317 SET last_committed_at_us = EXCLUDED.last_committed_at_us
2318 WHERE activity.last_committed_at_us < EXCLUDED.last_committed_at_us
2319 RETURNING key_id
2320 ) SELECT COUNT(*) FROM matched"
2321 ).bind(&keys).bind(&acct).bind(&at).fetch_one(&mut *tx).await.map_err(storage)?;
2322 report.unattributed = Some(expected_event_rows.checked_sub(
2323 u64::try_from(attributed).map_err(|_| StoreError("negative attribution count".into()))?
2324 ).ok_or_else(|| StoreError("attribution count exceeds accepted events".into()))?);
2325
2326 Ok(report)
2327 }
2328 .await;
2329 finish_transaction(tx, result).await
2333 }
2334}
2335
2336fn decode_status(stored: String) -> Result<AccountStatus, StoreError> {
2344 match stored.as_str() {
2345 s if s == AccountStatus::Active.as_str() => Ok(AccountStatus::Active),
2346 s if s == AccountStatus::Suspended.as_str() => Ok(AccountStatus::Suspended),
2347 s if s == AccountStatus::Closed.as_str() => Ok(AccountStatus::Closed),
2348 other => Err(StoreError(format!("unrecognized account status {other:?}"))),
2349 }
2350}
2351
2352fn decode_capacity_class(stored: String) -> Result<CapacityClass, StoreError> {
2361 match stored.as_str() {
2362 s if s == CapacityClass::Assured.as_str() => Ok(CapacityClass::Assured),
2363 s if s == CapacityClass::BestEffort.as_str() => Ok(CapacityClass::BestEffort),
2364 other => Err(StoreError(format!("unrecognized capacity class {other:?}"))),
2365 }
2366}
2367
2368fn decode_schedule(
2377 allowance: Option<i64>,
2378 period: Option<String>,
2379 rollover: Option<String>,
2380) -> Result<Option<BudgetSchedule>, StoreError> {
2381 let populated = [allowance.is_some(), period.is_some(), rollover.is_some()];
2382 let (Some(allowance), Some(period), Some(rollover)) = (allowance, period, rollover) else {
2383 if populated.iter().any(|present| *present) {
2384 return Err(StoreError(
2385 "stored budget schedule is partially populated".into(),
2386 ));
2387 }
2388 return Ok(None);
2389 };
2390 let period = match period.as_str() {
2391 s if s == Period::UtcCalendarMonth.as_str() => Period::UtcCalendarMonth,
2392 other => return Err(StoreError(format!("unrecognized budget period {other:?}"))),
2393 };
2394 let rollover = match rollover.as_str() {
2395 s if s == Rollover::None.as_str() => Rollover::None,
2396 other => {
2397 return Err(StoreError(format!(
2398 "unrecognized budget rollover {other:?}"
2399 )));
2400 }
2401 };
2402 Ok(Some(BudgetSchedule {
2403 allowance: to_units(allowance, "budget allowance")?,
2404 period,
2405 rollover,
2406 }))
2407}
2408
2409const BUDGET_VIEW_SQL: &str = "SELECT account_id, deposited, overage_recorded, usage_recorded,
2412 settlement_loss, expired, budget_allowance, budget_period,
2413 budget_rollover, period_start_us
2414 FROM tollgate_accounts WHERE account_id = $1";
2415
2416const GRANT_DEBIT_SQL: &str = "UPDATE tollgate_accounts
2419 SET balance = balance - $2,
2420 allowance_balance = allowance_balance - $3,
2421 next_fence = next_fence + 1
2422 WHERE account_id = $1
2423 RETURNING account_id, deposited, overage_recorded, usage_recorded,
2424 settlement_loss, expired, budget_allowance, budget_period,
2425 budget_rollover, period_start_us";
2426
2427fn budget_view(row: &sqlx::postgres::PgRow) -> Result<BudgetView, StoreError> {
2442 let deposited = to_units(row.get::<i64, _>(1), "deposited")?;
2443 let overage = to_units(row.get::<i64, _>(2), "overage_recorded")?;
2444 let usage = to_units(row.get::<i64, _>(3), "usage_recorded")?;
2445 let loss = to_units(row.get::<i64, _>(4), "settlement_loss")?;
2446 let expired = to_units(row.get::<i64, _>(5), "expired")?;
2447 let schedule = decode_schedule(row.get(6), row.get(7), row.get(8))?;
2448 let period_start = micros_ts(row.get::<i64, _>(9), "period_start_us")?;
2449
2450 let funded = deposited
2451 .checked_add(overage)
2452 .ok_or_else(|| StoreError("account funding total overflows".into()))?;
2453 let consumed = usage
2454 .checked_add(loss)
2455 .and_then(|spent| spent.checked_add(expired))
2456 .ok_or_else(|| StoreError("account consumption total overflows".into()))?;
2457 Ok(BudgetView {
2458 balance_at_publish: funded.checked_sub(consumed).ok_or_else(|| {
2459 StoreError(format!(
2460 "consumption {} exceeds funding {}",
2461 consumed.get(),
2462 funded.get()
2463 ))
2464 })?,
2465 period_end: schedule.map(|schedule| schedule.period.end_after(period_start)),
2470 })
2471}
2472
2473fn decode_publishable(
2485 principal: Principal,
2486 generation: i64,
2487 value: serde_json::Value,
2488) -> Result<PublishableSnapshot, StoreError> {
2489 let generation = generation_from(generation)?;
2490 let snapshot: StoredSnapshot =
2491 serde_json::from_value(value).map_err(|e| StoreError(format!("snapshot decode: {e}")))?;
2492 let budget = snapshot.budget;
2497 let publishable = PublishableSnapshot::try_new(Arc::new(snapshot.into_snapshot(generation)))
2498 .map_err(|error| {
2499 StoreError(format!(
2500 "invalid stored snapshot for principal {:#034x}: {error}",
2501 principal.0
2502 ))
2503 })?;
2504 Ok(match budget {
2505 Some(budget) => publishable.with_budget(Some(budget)),
2506 None => publishable,
2507 })
2508}
2509
2510fn generation_from(column: i64) -> Result<Generation, StoreError> {
2522 u64::try_from(column)
2523 .map(Generation)
2524 .map_err(|_| StoreError("stored snapshot generation is negative".into()))
2525}
2526
2527#[async_trait]
2528impl SnapshotSource for PostgresStore {
2529 async fn snapshot(&self, principal: Principal) -> Result<SnapshotResolution, StoreError> {
2530 let row = sqlx::query(
2531 "SELECT generation, snapshot, deleted FROM tollgate_snapshots WHERE principal = $1",
2532 )
2533 .bind(id_bytes(principal.0))
2534 .fetch_optional(&self.pool)
2535 .await
2536 .map_err(storage)?;
2537 match row {
2538 Some(row) if row.get::<bool, _>(2) => Ok(SnapshotResolution::Revoked {
2544 generation: generation_from(row.get::<i64, _>(0))?,
2545 }),
2546 Some(row) => Ok(SnapshotResolution::Present(decode_publishable(
2547 principal,
2548 row.get::<i64, _>(0),
2549 row.get(1),
2550 )?)),
2551 None => Ok(SnapshotResolution::Unknown),
2552 }
2553 }
2554
2555 fn subscribe(&self) -> broadcast::Receiver<SnapshotPush> {
2556 self.push.subscribe()
2557 }
2558
2559 async fn principals(&self) -> Result<Option<Vec<Principal>>, StoreError> {
2567 let rows = sqlx::query("SELECT principal FROM tollgate_snapshots ORDER BY principal")
2568 .fetch_all(&self.pool)
2569 .await
2570 .map_err(storage)?;
2571 Ok(Some(
2572 rows.iter()
2573 .map(|row| Principal(id_from(row.get::<Vec<u8>, _>(0).as_slice())))
2574 .collect(),
2575 ))
2576 }
2577}
2578
2579#[async_trait]
2580impl StoreHealth for PostgresStore {
2581 async fn ping(&self) -> Result<(), StoreError> {
2582 sqlx::query("SELECT 1")
2583 .execute(&self.pool)
2584 .await
2585 .map_err(storage)?;
2586 Ok(())
2587 }
2588}
2589
2590async fn republish_patched_snapshots(
2620 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
2621 account: AccountId,
2622 json_path: &'static str,
2623 value: &str,
2624) -> Result<(Vec<(Principal, PublishableSnapshot)>, usize), SetStatusError> {
2625 let rows = sqlx::query(
2626 "UPDATE tollgate_snapshots
2627 SET generation = generation + 1,
2628 snapshot = jsonb_set(snapshot, $3::text[], to_jsonb($2::text))
2629 WHERE account_id = $1
2630 AND deleted = FALSE
2631 AND snapshot #>> $3::text[] IS DISTINCT FROM $2::text
2632 RETURNING principal, generation, snapshot",
2633 )
2634 .bind(id_bytes(account.0))
2635 .bind(value)
2636 .bind(json_path)
2637 .fetch_all(&mut **tx)
2638 .await
2639 .map_err(storage)?;
2640
2641 let mut republished = Vec::with_capacity(rows.len());
2642 let mut unreadable = 0usize;
2643 for row in rows {
2644 let principal = Principal(id_from(row.get::<Vec<u8>, _>(0).as_slice()));
2645 match decode_publishable(
2646 principal,
2647 row.get::<i64, _>(1),
2648 row.get::<serde_json::Value, _>(2),
2649 ) {
2650 Ok(snapshot) => republished.push((principal, snapshot)),
2651 Err(error) => {
2659 unreadable += 1;
2660 tracing::warn!(
2661 %principal,
2662 %error,
2663 "restamped snapshot could not be decoded for push"
2664 );
2665 }
2666 }
2667 }
2668 republished.sort_unstable_by_key(|(principal, _)| *principal);
2671 Ok((republished, unreadable))
2672}
2673#[async_trait]
2674impl AdminStore for PostgresStore {
2675 async fn create_account(
2676 &self,
2677 config: AccountConfig,
2678 ) -> Result<tollgate_store::AdminReceipt<()>, CreateAccountError> {
2679 let result = sqlx::query(
2680 "INSERT INTO tollgate_accounts
2681 (account_id, balance, deposited, status, capacity_class, next_fence,
2682 usage_recorded, settlement_loss, overage_recorded)
2683 VALUES ($1, $2, $2, $3, $4, 1, 0, 0, 0)
2684 ON CONFLICT (account_id) DO NOTHING",
2685 )
2686 .bind(id_bytes(config.account_id.0))
2687 .bind(to_i64(config.initial_balance, "balance").map_err(CreateAccountError::Storage)?)
2688 .bind(config.status.as_str())
2689 .bind(config.capacity_class.as_str())
2690 .execute(&self.pool)
2691 .await
2692 .map_err(|e| CreateAccountError::Storage(storage(e)))?;
2693 if result.rows_affected() == 0 {
2694 return Err(CreateAccountError::AlreadyExists);
2695 }
2696 Ok(AdminReceipt::new(
2697 (),
2698 AdminState::Absent,
2699 AdminState::AccountCreated {
2700 initial_balance: config.initial_balance,
2701 status: config.status,
2702 capacity_class: config.capacity_class,
2703 },
2704 ))
2705 }
2706
2707 async fn deposit(
2708 &self,
2709 account: AccountId,
2710 units: CostUnits,
2711 ) -> Result<AdminReceipt<()>, AllocateError> {
2712 let row = sqlx::query(
2713 "UPDATE tollgate_accounts
2714 SET balance = balance + $2, deposited = deposited + $2
2715 WHERE account_id = $1
2716 RETURNING balance - $2 AS old_topup, deposited - $2 AS old_deposited,
2717 balance AS new_topup, deposited AS new_deposited",
2718 )
2719 .bind(id_bytes(account.0))
2720 .bind(i64::try_from(units.get()).map_err(|_| AllocateError::BalanceOverflow)?)
2721 .fetch_optional(&self.pool)
2722 .await
2723 .map_err(|error| {
2724 if error
2728 .as_database_error()
2729 .and_then(|db| db.code())
2730 .as_deref()
2731 == Some("22003")
2732 {
2733 AllocateError::BalanceOverflow
2734 } else {
2735 alloc_storage(error)
2736 }
2737 })?
2738 .ok_or(AllocateError::UnknownAccount)?;
2739 let state = |topup: &str, deposited: &str| -> Result<AdminState, StoreError> {
2740 Ok(AdminState::Funding {
2741 topup: to_units(row.get(topup), "audit topup")?,
2742 deposited: to_units(row.get(deposited), "audit deposited")?,
2743 })
2744 };
2745 Ok(AdminReceipt::new(
2746 (),
2747 state("old_topup", "old_deposited")?,
2748 state("new_topup", "new_deposited")?,
2749 ))
2750 }
2751
2752 async fn set_budget_schedule(
2753 &self,
2754 account: AccountId,
2755 schedule: Option<BudgetSchedule>,
2756 ) -> Result<AdminReceipt<()>, BudgetError> {
2757 let allowance = schedule
2767 .map(|s| to_i64(s.allowance, "budget allowance"))
2768 .transpose()
2769 .map_err(BudgetError::Storage)?;
2770 let mut tx = self.pool.begin().await.map_err(storage)?;
2771 let result = async {
2772 let row = sqlx::query(
2776 "SELECT budget_allowance, budget_period, budget_rollover
2777 FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE",
2778 )
2779 .bind(id_bytes(account.0))
2780 .fetch_optional(&mut *tx)
2781 .await
2782 .map_err(storage)?
2783 .ok_or(BudgetError::UnknownAccount)?;
2784 let before = decode_schedule(row.get(0), row.get(1), row.get(2))?;
2785 sqlx::query(
2786 "UPDATE tollgate_accounts
2787 SET budget_allowance = $2, budget_period = $3, budget_rollover = $4
2788 WHERE account_id = $1",
2789 )
2790 .bind(id_bytes(account.0))
2791 .bind(allowance)
2792 .bind(schedule.map(|s| s.period.as_str()))
2793 .bind(schedule.map(|s| s.rollover.as_str()))
2794 .execute(&mut *tx)
2795 .await
2796 .map_err(storage)?;
2797 Ok(AdminReceipt::new(
2798 (),
2799 AdminState::Budget { schedule: before },
2800 AdminState::Budget { schedule },
2801 ))
2802 }
2803 .await;
2804 finish_transaction(tx, result).await
2805 }
2806
2807 async fn account_view(&self, account: AccountId) -> Result<Option<AccountView>, StoreError> {
2808 let mut tx = self.pool.begin().await.map_err(storage)?;
2817 let result = async {
2818 sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY")
2819 .execute(&mut *tx)
2820 .await
2821 .map_err(storage)?;
2822 let account_row = sqlx::query(
2823 "SELECT deposited, balance, usage_recorded, settlement_loss, overage_recorded,
2824 expired, status, capacity_class, budget_allowance, budget_period,
2825 budget_rollover, period_start_us
2826 FROM tollgate_accounts WHERE account_id = $1",
2827 )
2828 .bind(id_bytes(account.0))
2829 .fetch_optional(&mut *tx)
2830 .await
2831 .map_err(storage)?;
2832 let lease_row = sqlx::query(ACTIVE_LEASE_SUM_SQL)
2833 .bind(id_bytes(account.0))
2834 .fetch_one(&mut *tx)
2835 .await
2836 .map_err(storage)?;
2837 Ok::<_, StoreError>((account_row, lease_row))
2838 }
2839 .await;
2840 let (account_row, lease_row) = finish_transaction(tx, result).await?;
2841
2842 let Some(row) = account_row else {
2843 return Ok(None);
2844 };
2845 let active_grants = to_units(lease_row.get::<i64, _>(0), "active lease grants")?;
2846 let active_used = to_units(lease_row.get::<i64, _>(1), "active lease usage")?;
2847 let recorded = to_units(row.get::<i64, _>(2), "usage_recorded")?;
2848 Ok(Some(AccountView {
2849 account_id: account,
2850 status: decode_status(row.get::<String, _>(6))?,
2851 capacity_class: decode_capacity_class(row.get::<String, _>(7))?,
2852 schedule: decode_schedule(
2853 row.get::<Option<i64>, _>(8),
2854 row.get::<Option<String>, _>(9),
2855 row.get::<Option<String>, _>(10),
2856 )?,
2857 period_start: StoredInstant {
2858 micros: row.get::<i64, _>(11),
2859 submicro_nanos: 0,
2860 }
2861 .timestamp()?,
2862 conservation: Conservation {
2863 deposited: to_units(row.get::<i64, _>(0), "deposited")?,
2864 overage_recorded: to_units(row.get::<i64, _>(4), "overage_recorded")?,
2865 balance: to_units(row.get::<i64, _>(1), "balance")?,
2866 active_lease_grants: active_grants,
2867 settled_usage: recorded.checked_sub(active_used).ok_or_else(|| {
2871 StoreError(format!(
2872 "active lease usage {} exceeds recorded usage {} for account {account}",
2873 active_used.get(),
2874 recorded.get()
2875 ))
2876 })?,
2877 settlement_loss: to_units(row.get::<i64, _>(3), "settlement_loss")?,
2878 expired: to_units(row.get::<i64, _>(5), "expired")?,
2879 },
2880 }))
2881 }
2882
2883 async fn roll_due_periods(
2884 &self,
2885 now: Timestamp,
2886 limit: NonZeroUsize,
2887 ) -> Result<RolloverBatch, StoreError> {
2888 let limit_i = i64::try_from(limit.get())
2889 .map_err(|_| StoreError(format!("rollover batch limit exceeds i64 range: {limit}")))?;
2890 let mut rolled = Vec::new();
2891 for period in Period::ALL {
2898 let boundary_us = ts_micros(period.start_of(now));
2899 let remaining = limit_i
2900 - i64::try_from(rolled.len())
2901 .map_err(|_| StoreError("rollover batch row count exceeds i64 range".into()))?;
2902 if remaining <= 0 {
2903 break;
2904 }
2905 let rows = sqlx::query(DUE_PERIODS_SQL)
2916 .bind(period.as_str())
2917 .bind(boundary_us)
2918 .bind(remaining)
2919 .fetch_all(&self.pool)
2920 .await
2921 .map_err(storage)?;
2922
2923 for row in rows {
2924 rolled.push((
2925 row.get::<i64, _>(3),
2926 RolledAccount {
2927 account_id: AccountId(id_from(&row.get::<Vec<u8>, _>(0))),
2928 deposited: to_units(row.get::<i64, _>(1), "budget allowance")?,
2929 expired: to_units(row.get::<i64, _>(2), "expiring allowance")?,
2930 },
2931 ));
2932 }
2933 }
2934 rolled
2941 .sort_unstable_by_key(|(crossed_from, account)| (*crossed_from, account.account_id.0));
2942 RolloverBatch::try_new(
2943 rolled.into_iter().map(|(_, account)| account).collect(),
2944 limit,
2945 )
2946 }
2947
2948 async fn set_account_status(
2949 &self,
2950 account: AccountId,
2951 status: AccountStatus,
2952 ) -> Result<tollgate_store::AdminReceipt<StatusChange>, SetStatusError> {
2953 let mut tx = self.pool.begin().await.map_err(storage)?;
2954 let result = async {
2955 let row = sqlx::query(
2960 "SELECT status FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE",
2961 )
2962 .bind(id_bytes(account.0))
2963 .fetch_optional(&mut *tx)
2964 .await
2965 .map_err(storage)?
2966 .ok_or(SetStatusError::UnknownAccount)?;
2967
2968 let before = decode_status(row.get::<String, _>(0))?;
2969 if before == AccountStatus::Closed && status != AccountStatus::Closed {
2970 return Err(SetStatusError::AccountClosed);
2973 }
2974
2975 sqlx::query("UPDATE tollgate_accounts SET status = $2 WHERE account_id = $1")
2976 .bind(id_bytes(account.0))
2977 .bind(status.as_str())
2978 .execute(&mut *tx)
2979 .await
2980 .map_err(storage)?;
2981
2982 let (republished, unreadable) =
2983 republish_patched_snapshots(&mut tx, account, "{status}", status.as_str()).await?;
2984 Ok((before, republished, unreadable))
2985 }
2986 .await;
2987
2988 let (before, republished, unreadable) = finish_transaction(tx, result).await?;
2989 if pushes_exceed_capacity(republished.len()) {
2991 tracing::warn!(
2992 %account,
2993 principals = republished.len(),
2994 capacity = PUSH_CHANNEL_CAPACITY,
2995 "status change emitted more pushes than the channel holds; subscribers will resync"
2996 );
2997 }
2998 let count = republished.len();
2999 for (principal, snapshot) in republished {
3000 self.push_to_subscribers(SnapshotPush {
3001 principal,
3002 resolution: SnapshotResolution::Present(snapshot),
3003 });
3004 }
3005 Ok(AdminReceipt::new(
3006 StatusChange {
3007 republished: count + unreadable,
3010 unreadable,
3011 },
3012 AdminState::Status { status: before },
3013 AdminState::Status { status },
3014 ))
3015 }
3016
3017 async fn set_capacity_class(
3018 &self,
3019 account: AccountId,
3020 class: CapacityClass,
3021 ) -> Result<tollgate_store::AdminReceipt<StatusChange>, SetStatusError> {
3022 let mut tx = self.pool.begin().await.map_err(storage)?;
3023 let result = async {
3024 let row = sqlx::query(
3032 "SELECT status, capacity_class FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE",
3033 )
3034 .bind(id_bytes(account.0))
3035 .fetch_optional(&mut *tx)
3036 .await
3037 .map_err(storage)?
3038 .ok_or(SetStatusError::UnknownAccount)?;
3039
3040 let before = decode_capacity_class(row.get::<String, _>(1))?;
3041 if decode_status(row.get::<String, _>(0))? == AccountStatus::Closed {
3046 return Err(SetStatusError::AccountClosed);
3047 }
3048
3049 sqlx::query("UPDATE tollgate_accounts SET capacity_class = $2 WHERE account_id = $1")
3050 .bind(id_bytes(account.0))
3051 .bind(class.as_str())
3052 .execute(&mut *tx)
3053 .await
3054 .map_err(storage)?;
3055
3056 let (republished, unreadable) =
3057 republish_patched_snapshots(&mut tx, account, "{capacity_class}", class.as_str()).await?;
3058 Ok((before, republished, unreadable))
3059 }
3060 .await;
3061
3062 let (before, republished, unreadable) = finish_transaction(tx, result).await?;
3063 if pushes_exceed_capacity(republished.len()) {
3064 tracing::warn!(
3065 %account,
3066 principals = republished.len(),
3067 capacity = PUSH_CHANNEL_CAPACITY,
3068 "capacity class change emitted more pushes than the channel holds; \
3069 subscribers will resync"
3070 );
3071 }
3072 let count = republished.len();
3073 for (principal, snapshot) in republished {
3074 self.push_to_subscribers(SnapshotPush {
3075 principal,
3076 resolution: SnapshotResolution::Present(snapshot),
3077 });
3078 }
3079 Ok(AdminReceipt::new(
3080 StatusChange {
3081 republished: count + unreadable,
3084 unreadable,
3085 },
3086 AdminState::CapacityClass {
3087 capacity_class: before,
3088 },
3089 AdminState::CapacityClass {
3090 capacity_class: class,
3091 },
3092 ))
3093 }
3094
3095 async fn publish_snapshot(
3096 &self,
3097 principal: Principal,
3098 snapshot: PublishableSnapshot,
3099 ) -> Result<tollgate_store::AdminReceipt<()>, PublishSnapshotError> {
3100 let generation = i64::try_from(snapshot.generation.0).map_err(|_| {
3101 StoreError("snapshot generation exceeds PostgreSQL BIGINT range".into())
3102 })?;
3103
3104 let mut tx = self.pool.begin().await.map_err(storage)?;
3105 let result = publish_in_tx(&mut tx, principal, generation, snapshot).await;
3106
3107 let (written, published, before, after) = finish_transaction(tx, result).await?;
3108 if written {
3109 self.push_to_subscribers(SnapshotPush {
3113 principal,
3114 resolution: SnapshotResolution::Present(published),
3115 });
3116 }
3117 Ok(AdminReceipt::new((), before, after))
3118 }
3119
3120 async fn remove_snapshot(&self, principal: Principal) -> Result<AdminReceipt<()>, StoreError> {
3121 let mut tx = self.pool.begin().await.map_err(storage)?;
3122 let result = remove_in_tx(&mut tx, principal).await;
3123 let receipt = finish_transaction(tx, result).await?;
3124 self.announce_removal(principal, &receipt);
3125 Ok(receipt)
3126 }
3127}
3128
3129async fn lock_account_key(
3153 tx: &mut Transaction<'_, Postgres>,
3154 account: AccountId,
3155 key: KeyId,
3156) -> Result<(Principal, bool), KeySnapshotError> {
3157 let row = sqlx::query(
3158 "SELECT principal, revoked_at_us IS NOT NULL
3159 FROM tollgate_credential_keys WHERE key_id = $1 AND account_id = $2 FOR SHARE",
3160 )
3161 .bind(id_bytes(key.0))
3162 .bind(id_bytes(account.0))
3163 .fetch_optional(&mut **tx)
3164 .await
3165 .map_err(storage)?
3166 .ok_or(KeySnapshotError::UnknownCredential)?;
3167 let principal: Vec<u8> = row.get(0);
3168 let principal: [u8; 16] = principal
3169 .try_into()
3170 .map_err(|_| StoreError("credential principal is not 16 bytes".into()))?;
3171 Ok((Principal(u128::from_be_bytes(principal)), row.get(1)))
3172}
3173
3174async fn publish_in_tx(
3180 tx: &mut Transaction<'_, Postgres>,
3181 principal: Principal,
3182 generation: i64,
3183 snapshot: PublishableSnapshot,
3184) -> Result<(bool, PublishableSnapshot, AdminState, AdminState), PublishSnapshotError> {
3185 if let Some(key_id) = snapshot.key_id {
3186 let matches: bool = sqlx::query_scalar(
3187 "SELECT EXISTS(SELECT 1 FROM tollgate_credential_keys
3188 WHERE key_id = $1 AND principal = $2 AND account_id = $3)",
3189 )
3190 .bind(id_bytes(key_id.0))
3191 .bind(id_bytes(principal.0))
3192 .bind(id_bytes(snapshot.account_id.0))
3193 .fetch_one(&mut **tx)
3194 .await
3195 .map_err(storage)?;
3196 if !matches {
3197 return Err(PublishSnapshotError::CredentialMismatch { key_id });
3198 }
3199 }
3200 let ledger = sqlx::query(
3212 "SELECT status, deposited, overage_recorded, usage_recorded, settlement_loss,
3213 expired, budget_allowance, budget_period, budget_rollover, period_start_us,
3214 capacity_class
3215 FROM tollgate_accounts WHERE account_id = $1 FOR SHARE",
3216 )
3217 .bind(id_bytes(snapshot.account_id.0))
3218 .fetch_optional(&mut **tx)
3219 .await
3220 .map_err(storage)?;
3221
3222 let view = match &ledger {
3230 Some(row) => {
3231 let ledger = decode_status(row.get::<String, _>(0))?;
3232 if ledger != snapshot.status {
3233 return Err(PublishSnapshotError::StatusMismatch {
3234 ledger,
3235 submitted: snapshot.status,
3236 });
3237 }
3238 let ledger_class = decode_capacity_class(row.get::<String, _>(10))?;
3242 if ledger_class != snapshot.capacity_class {
3243 return Err(PublishSnapshotError::CapacityClassMismatch {
3244 ledger: ledger_class,
3245 submitted: snapshot.capacity_class,
3246 });
3247 }
3248 Some(budget_view(row)?)
3249 }
3250 None => None,
3251 };
3252 let published = snapshot.with_budget(view);
3253 let value = serde_json::to_value(StoredSnapshotRef::from(published.as_snapshot()))
3254 .map_err(|e| StoreError(format!("snapshot encode: {e}")))?;
3255
3256 let (written, before, after) = write_snapshot_audited(tx, principal, generation, value).await?;
3257 Ok((written, published, before, after))
3258}
3259
3260async fn remove_in_tx(
3263 tx: &mut Transaction<'_, Postgres>,
3264 principal: Principal,
3265) -> Result<AdminReceipt<()>, StoreError> {
3266 let before = snapshot_audit_row(tx, principal).await?;
3267 let after = match before {
3268 AdminState::Snapshot {
3269 generation,
3270 revoked: false,
3271 } => {
3272 sqlx::query("UPDATE tollgate_snapshots SET deleted = TRUE WHERE principal = $1")
3273 .bind(id_bytes(principal.0))
3274 .execute(&mut **tx)
3275 .await
3276 .map_err(storage)?;
3277 AdminState::Snapshot {
3278 generation,
3279 revoked: true,
3280 }
3281 }
3282 state => state,
3283 };
3284 Ok(AdminReceipt::new((), before, after))
3285}
3286
3287async fn snapshot_audit_row(
3290 tx: &mut Transaction<'_, Postgres>,
3291 principal: Principal,
3292) -> Result<AdminState, StoreError> {
3293 let row = sqlx::query(
3294 "SELECT generation, deleted FROM tollgate_snapshots WHERE principal = $1 FOR UPDATE",
3295 )
3296 .bind(id_bytes(principal.0))
3297 .fetch_optional(&mut **tx)
3298 .await
3299 .map_err(storage)?;
3300 row.map(|row| {
3301 Ok(AdminState::Snapshot {
3302 generation: generation_from(row.get(0))?,
3303 revoked: row.get(1),
3304 })
3305 })
3306 .unwrap_or(Ok(AdminState::Absent))
3307}
3308
3309async fn write_snapshot_audited(
3310 tx: &mut Transaction<'_, Postgres>,
3311 principal: Principal,
3312 generation: i64,
3313 value: serde_json::Value,
3314) -> Result<(bool, AdminState, AdminState), StoreError> {
3315 let mut before = snapshot_audit_row(tx, principal).await?;
3316 let after = AdminState::Snapshot {
3317 generation: generation_from(generation)?,
3318 revoked: false,
3319 };
3320 if before == AdminState::Absent {
3321 let inserted = sqlx::query(
3322 "INSERT INTO tollgate_snapshots (principal, generation, snapshot, deleted)
3323 VALUES ($1, $2, $3, FALSE) ON CONFLICT (principal) DO NOTHING",
3324 )
3325 .bind(id_bytes(principal.0))
3326 .bind(generation)
3327 .bind(&value)
3328 .execute(&mut **tx)
3329 .await
3330 .map_err(storage)?;
3331 if inserted.rows_affected() == 1 {
3332 return Ok((true, before, after));
3333 }
3334 before = snapshot_audit_row(tx, principal).await?;
3337 }
3338 let AdminState::Snapshot {
3339 generation: previous,
3340 ..
3341 } = before
3342 else {
3343 return Err(StoreError("snapshot disappeared during publication".into()));
3344 };
3345 if previous >= generation_from(generation)? {
3346 return Ok((false, before, before));
3347 }
3348 sqlx::query("UPDATE tollgate_snapshots SET generation = $2, snapshot = $3, deleted = FALSE WHERE principal = $1")
3349 .bind(id_bytes(principal.0)).bind(generation).bind(value)
3350 .execute(&mut **tx).await.map_err(storage)?;
3351 Ok((true, before, after))
3352}
3353
3354#[cfg(test)]
3355mod tests {
3356 use super::*;
3357
3358 #[test]
3359 fn stored_fences_use_the_exact_positive_bigint_domain() {
3360 for invalid in [i64::MIN, -1, 0] {
3361 assert!(stored_fence(invalid).is_err());
3362 }
3363 for valid in [1, 2, i64::MAX] {
3364 assert_eq!(stored_fence(valid).unwrap(), FencingToken(valid as u64));
3365 }
3366 }
3367
3368 #[tokio::test]
3371 async fn ping_surfaces_a_closed_pool() {
3372 let pool = PgPoolOptions::new()
3373 .connect_lazy("postgres://localhost/tollgate")
3374 .unwrap();
3375 pool.close().await;
3376 let (push, _) = broadcast::channel(1);
3377 let store = PostgresStore {
3378 pool,
3379 policy: GrantPolicy::default(),
3380 push,
3381 };
3382
3383 assert!(store.ping().await.is_err());
3384 assert!(matches!(
3385 AdminStore::deposit(&store, AccountId(1), CostUnits(1)).await,
3386 Err(AllocateError::Storage(_))
3387 ));
3388 }
3389}