1use crate::ports::{Database, DbError, Statement};
21use std::collections::HashMap;
22use std::time::Duration as StdDuration;
23use time::{Date, Month, OffsetDateTime, UtcOffset};
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub struct Months(pub u32);
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum Period {
34 CalendarMonthUtc,
36 Anchored {
44 anchor: OffsetDateTime,
45 every: Months,
46 },
47 Day,
49}
50
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub struct PeriodWindow {
54 pub start: OffsetDateTime,
55 pub end: OffsetDateTime,
56}
57
58impl Period {
59 #[must_use]
61 pub fn window_at(self, now: OffsetDateTime) -> PeriodWindow {
62 self.window_for(self.index_at(now))
63 }
64
65 #[must_use]
68 pub fn index_at(self, now: OffsetDateTime) -> i64 {
69 let now = now.to_offset(UtcOffset::UTC);
70 match self {
71 Self::CalendarMonthUtc => month_index(now.year(), now.month()),
72 Self::Day => i64::from(now.date().to_julian_day()),
73 Self::Anchored { anchor, every } => {
74 let anchor = anchor.to_offset(UtcOffset::UTC);
75 let every = i64::from(every.0.max(1));
76 let diff = month_index(now.year(), now.month())
79 - month_index(anchor.year(), anchor.month());
80 let mut index = diff.div_euclid(every);
81 let (lowest, highest) = month_index_bounds();
86 while index < highest && anchored_start(anchor, every, index + 1) <= now {
87 index += 1;
88 }
89 while index > lowest && anchored_start(anchor, every, index) > now {
90 index -= 1;
91 }
92 index
93 }
94 }
95 }
96
97 #[must_use]
101 pub fn window_for(self, index: i64) -> PeriodWindow {
102 let next = index.saturating_add(1);
103 match self {
104 Self::CalendarMonthUtc => PeriodWindow {
105 start: month_start(index),
106 end: month_start(next),
107 },
108 Self::Day => PeriodWindow {
109 start: day_start(index),
110 end: day_start(next),
111 },
112 Self::Anchored { anchor, every } => {
113 let every = i64::from(every.0.max(1));
114 PeriodWindow {
115 start: anchored_start(anchor, every, index),
116 end: anchored_start(anchor, every, next),
117 }
118 }
119 }
120 }
121}
122
123fn month_index(year: i32, month: Month) -> i64 {
126 i64::from(year) * 12 + i64::from(u8::from(month)) - 1
127}
128
129fn month_index_bounds() -> (i64, i64) {
139 (
140 month_index(Date::MIN.year(), Month::January),
141 month_index(Date::MAX.year(), Month::December),
142 )
143}
144
145fn clamp_month_index(index: i64) -> i64 {
147 let (lowest, highest) = month_index_bounds();
148 index.clamp(lowest, highest)
149}
150
151fn month_start(index: i64) -> OffsetDateTime {
154 let index = clamp_month_index(index);
155 let year = i32::try_from(index.div_euclid(12)).expect("a clamped month index fits an i32 year");
156 let month = month_from(index.rem_euclid(12) + 1);
157 Date::from_calendar_date(year, month, 1)
158 .expect("the first of a valid month is a valid date")
159 .midnight()
160 .assume_utc()
161}
162
163fn day_start(index: i64) -> OffsetDateTime {
166 let lowest = i64::from(Date::MIN.to_julian_day());
167 let highest = i64::from(Date::MAX.to_julian_day());
168 let julian =
169 i32::try_from(index.clamp(lowest, highest)).expect("a clamped julian day fits i32");
170 Date::from_julian_day(julian)
171 .expect("a valid julian day is a valid date")
172 .midnight()
173 .assume_utc()
174}
175
176fn anchored_start(anchor: OffsetDateTime, every: i64, index: i64) -> OffsetDateTime {
182 let anchor = anchor.to_offset(UtcOffset::UTC);
183 let target = clamp_month_index(
184 month_index(anchor.year(), anchor.month()).saturating_add(index.saturating_mul(every)),
185 );
186 let year =
187 i32::try_from(target.div_euclid(12)).expect("a clamped month index fits an i32 year");
188 let month = month_from(target.rem_euclid(12) + 1);
189 let day = anchor.day().min(days_in_month(year, month));
190 Date::from_calendar_date(year, month, day)
191 .expect("the clamped day is valid for its month")
192 .with_hms(anchor.hour(), anchor.minute(), anchor.second())
193 .expect("the anchor's time-of-day is valid")
194 .assume_utc()
195}
196
197fn month_from(one_based: i64) -> Month {
199 let number = u8::try_from(one_based).expect("1..=12 fits u8");
200 Month::try_from(number).expect("1..=12 is a month")
201}
202
203fn days_in_month(year: i32, month: Month) -> u8 {
206 match month {
207 Month::February if is_leap_year(year) => 29,
208 Month::February => 28,
209 Month::April | Month::June | Month::September | Month::November => 30,
210 _ => 31,
211 }
212}
213
214fn is_leap_year(year: i32) -> bool {
216 (year % 4 == 0 && year % 100 != 0) || year % 400 == 0
217}
218
219pub(crate) fn rfc3339(at: OffsetDateTime) -> String {
223 use time::format_description::well_known::Rfc3339;
224 let at = at.to_offset(UtcOffset::UTC);
225 at.replace_nanosecond(0)
226 .unwrap_or(at)
227 .format(&Rfc3339)
228 .unwrap_or_default()
229}
230
231fn counter(value: u64) -> i64 {
235 i64::try_from(value).unwrap_or(i64::MAX)
236}
237
238fn retry_after(period_end: OffsetDateTime, now: OffsetDateTime) -> StdDuration {
241 let millis = (period_end - now).whole_milliseconds().max(0);
242 StdDuration::from_millis(u64::try_from(millis).unwrap_or(u64::MAX))
243}
244
245#[derive(Debug, Clone, Copy, PartialEq, Eq)]
247pub struct Consumed {
248 pub used: u64,
250 pub limit: Option<u64>,
252 pub period_end: OffsetDateTime,
254}
255
256#[derive(Debug, Clone, Copy, PartialEq, Eq)]
261pub struct Exhausted {
262 pub used: u64,
264 pub limit: u64,
266 pub period_end: OffsetDateTime,
268 pub retry_after: StdDuration,
271}
272
273impl Exhausted {
274 #[must_use]
279 pub fn spent(used: u64, limit: u64, period: Period, now: OffsetDateTime) -> Self {
280 let end = period.window_at(now).end;
281 Self {
282 used,
283 limit,
284 period_end: end,
285 retry_after: retry_after(end, now),
286 }
287 }
288}
289
290#[derive(Debug, Clone, Copy, PartialEq, Eq)]
292pub enum Consumption {
293 Consumed(Consumed),
294 Exhausted(Exhausted),
295}
296
297#[derive(Debug, Clone)]
304pub struct Usage {
305 table: String,
306}
307
308impl Usage {
309 #[must_use]
310 pub fn new(table: impl Into<String>) -> Self {
311 Self {
312 table: table.into(),
313 }
314 }
315
316 #[must_use]
319 pub fn create_table_sql(&self) -> String {
320 format!(
321 "CREATE TABLE IF NOT EXISTS {table} (\n \
322 subject TEXT NOT NULL,\n \
323 meter TEXT NOT NULL,\n \
324 period_start TEXT NOT NULL,\n \
325 used BIGINT NOT NULL,\n \
326 PRIMARY KEY (subject, meter, period_start)\n);",
327 table = self.table
328 )
329 }
330
331 #[must_use]
337 pub fn consume_statement(
338 &self,
339 subject: &str,
340 meter: &str,
341 period: Period,
342 now: OffsetDateTime,
343 amount: u64,
344 limit: Option<u64>,
345 ) -> Statement {
346 let start = rfc3339(period.window_at(now).start);
347 match limit {
348 Some(limit) => Statement::with_values(
349 guard_sql(&self.table),
350 vec![
351 subject.to_owned().into(),
352 meter.to_owned().into(),
353 start.into(),
354 counter(amount).into(),
355 counter(limit).into(),
356 counter(amount).into(),
357 counter(limit).into(),
358 ],
359 ),
360 None => Statement::with_values(
361 plain_sql(&self.table),
362 vec![
363 subject.to_owned().into(),
364 meter.to_owned().into(),
365 start.into(),
366 counter(amount).into(),
367 ],
368 ),
369 }
370 }
371
372 #[allow(clippy::too_many_arguments)]
386 pub async fn consume(
387 &self,
388 db: &dyn Database,
389 subject: &str,
390 meter: &str,
391 period: Period,
392 now: OffsetDateTime,
393 amount: u64,
394 limit: Option<u64>,
395 ) -> Result<Consumption, DbError> {
396 let window = period.window_at(now);
397 let start = rfc3339(window.start);
398 let (sql, values) = match limit {
402 Some(limit) => (
403 guarded_sql(&self.table),
404 vec![
405 subject.to_owned().into(),
406 meter.to_owned().into(),
407 start.into(),
408 counter(amount).into(),
409 counter(amount).into(),
410 counter(limit).into(),
411 counter(limit).into(),
412 ],
413 ),
414 None => (
415 plain_sql(&self.table),
416 vec![
417 subject.to_owned().into(),
418 meter.to_owned().into(),
419 start.into(),
420 counter(amount).into(),
421 ],
422 ),
423 };
424 let affected = db.execute(&Statement::with_values(sql, values)).await?;
425 let used = self.read(db, subject, meter, period, now).await?;
426 Ok(match limit {
427 Some(limit) if affected == 0 => Consumption::Exhausted(Exhausted {
428 used,
429 limit,
430 period_end: window.end,
431 retry_after: retry_after(window.end, now),
432 }),
433 _ => Consumption::Consumed(Consumed {
434 used,
435 limit,
436 period_end: window.end,
437 }),
438 })
439 }
440
441 pub async fn refund(
448 &self,
449 db: &dyn Database,
450 subject: &str,
451 meter: &str,
452 period: Period,
453 now: OffsetDateTime,
454 amount: u64,
455 ) -> Result<(), DbError> {
456 let start = rfc3339(period.window_at(now).start);
457 let stmt = Statement::with_values(
460 format!(
461 "UPDATE {table} SET used = CASE WHEN used > ? THEN used - ? ELSE 0 END\n\
462 WHERE subject = ? AND meter = ? AND period_start = ?",
463 table = self.table
464 ),
465 vec![
466 counter(amount).into(),
467 counter(amount).into(),
468 subject.to_owned().into(),
469 meter.to_owned().into(),
470 start.into(),
471 ],
472 );
473 db.execute(&stmt).await?;
474 Ok(())
475 }
476
477 pub async fn read(
484 &self,
485 db: &dyn Database,
486 subject: &str,
487 meter: &str,
488 period: Period,
489 now: OffsetDateTime,
490 ) -> Result<u64, DbError> {
491 let start = rfc3339(period.window_at(now).start);
492 let rows = db
493 .query(&Statement::with_values(
494 format!(
495 "SELECT used FROM {table} WHERE subject = ? AND meter = ? AND period_start = ?",
496 table = self.table
497 ),
498 vec![
499 subject.to_owned().into(),
500 meter.to_owned().into(),
501 start.into(),
502 ],
503 ))
504 .await?;
505 Ok(rows
506 .first()
507 .and_then(|row| row.get::<u64>("used"))
508 .unwrap_or(0))
509 }
510
511 pub async fn history(
519 &self,
520 db: &dyn Database,
521 subject: &str,
522 meter: &str,
523 period: Period,
524 now: OffsetDateTime,
525 last_n_periods: u32,
526 ) -> Result<Vec<(PeriodWindow, u64)>, DbError> {
527 let n = i64::from(last_n_periods);
528 if n == 0 {
529 return Ok(Vec::new());
530 }
531 let current = period.index_at(now);
532 let oldest = period.window_for(current - n + 1).start;
533 let newest_end = period.window_for(current).end;
534 let rows = db
535 .query(&Statement::with_values(
536 format!(
537 "SELECT period_start, used FROM {table} \
538 WHERE subject = ? AND meter = ? AND period_start >= ? AND period_start < ?",
539 table = self.table
540 ),
541 vec![
542 subject.to_owned().into(),
543 meter.to_owned().into(),
544 rfc3339(oldest).into(),
545 rfc3339(newest_end).into(),
546 ],
547 ))
548 .await?;
549 let mut spent: HashMap<String, u64> = HashMap::new();
550 for row in &rows.rows {
551 if let (Some(start), Some(used)) =
552 (row.get::<String>("period_start"), row.get::<u64>("used"))
553 {
554 spent.insert(start, used);
555 }
556 }
557 Ok((0..n)
558 .map(|back| {
559 let window = period.window_for(current - back);
560 let used = spent.get(&rfc3339(window.start)).copied().unwrap_or(0);
561 (window, used)
562 })
563 .collect())
564 }
565
566 pub async fn purge(
579 &self,
580 db: &dyn Database,
581 period: Period,
582 now: OffsetDateTime,
583 keep_periods: u32,
584 ) -> Result<u64, DbError> {
585 let oldest_kept = period.index_at(now) - i64::from(keep_periods) + 1;
586 let cutoff = rfc3339(period.window_for(oldest_kept).start);
587 db.execute(&Statement::with_values(
588 format!(
589 "DELETE FROM {table} WHERE period_start < ?",
590 table = self.table
591 ),
592 vec![cutoff.into()],
593 ))
594 .await
595 }
596}
597
598fn guard_sql(table: &str) -> String {
602 format!(
603 "INSERT INTO {table} (subject, meter, period_start, used)\n\
604 VALUES (?, ?, ?, CASE WHEN ? <= ? THEN ? ELSE NULL END)\n\
605 ON CONFLICT (subject, meter, period_start) DO UPDATE\n\
606 SET used = CASE WHEN {table}.used + excluded.used <= ? \
607 THEN {table}.used + excluded.used ELSE NULL END"
608 )
609}
610
611fn plain_sql(table: &str) -> String {
613 format!(
614 "INSERT INTO {table} (subject, meter, period_start, used)\n\
615 VALUES (?, ?, ?, ?)\n\
616 ON CONFLICT (subject, meter, period_start) DO UPDATE\n\
617 SET used = {table}.used + excluded.used"
618 )
619}
620
621fn guarded_sql(table: &str) -> String {
626 format!(
627 "INSERT INTO {table} (subject, meter, period_start, used)\n\
628 SELECT ?, ?, ?, ? WHERE ? <= ?\n\
629 ON CONFLICT (subject, meter, period_start) DO UPDATE\n\
630 SET used = {table}.used + excluded.used\n\
631 WHERE {table}.used + excluded.used <= ?"
632 )
633}
634
635#[cfg(test)]
636mod tests {
637 use super::*;
638 use crate::lint_portable_sql;
639 use time::format_description::well_known::Rfc3339;
640
641 fn at(text: &str) -> OffsetDateTime {
644 OffsetDateTime::parse(text, &Rfc3339).expect("test instant parses")
645 }
646
647 fn window(start: &str, end: &str) -> PeriodWindow {
648 PeriodWindow {
649 start: at(start),
650 end: at(end),
651 }
652 }
653
654 #[test]
655 fn create_table_sql_is_portable_ddl() {
656 let sql = Usage::new("waitlist_usage").create_table_sql();
657 assert!(sql.contains("used BIGINT NOT NULL"));
658 assert_eq!(lint_portable_sql(&sql), vec![], "DDL must be portable");
659 }
660
661 #[test]
662 fn the_statement_shapes_bind_seven_or_four_values() {
663 let usage = Usage::new("waitlist_usage");
664 let period = Period::CalendarMonthUtc;
665 let now = at("2026-01-15T00:00:00Z");
666
667 let guarded = usage.consume_statement("a", "mail", period, now, 1, Some(100));
668 assert!(guarded.sql.contains("ELSE NULL END"));
669 assert!(guarded.sql.contains("excluded.used"));
670 assert_eq!(guarded.values.0.len(), 7);
672
673 let unbounded = usage.consume_statement("a", "mail", period, now, 1, None);
674 assert_eq!(unbounded.values.0.len(), 4);
675
676 let huge = usage.consume_statement("a", "mail", period, now, u64::MAX, Some(u64::MAX));
678 assert!(
679 matches!(huge.values.0[3], sea_query::Value::BigInt(Some(v)) if v == i64::MAX),
680 "amount must saturate to i64::MAX"
681 );
682 }
683
684 #[test]
685 fn calendar_and_day_windows_are_utc() {
686 assert_eq!(
688 Period::CalendarMonthUtc.window_at(at("2026-12-15T00:00:00Z")),
689 window("2026-12-01T00:00:00Z", "2027-01-01T00:00:00Z")
690 );
691 assert_eq!(
693 Period::Day.window_at(at("2026-03-01T23:59:59Z")),
694 window("2026-03-01T00:00:00Z", "2026-03-02T00:00:00Z")
695 );
696 let local = at("2026-03-01T00:30:00+05:30");
699 assert_eq!(
700 Period::Day.window_at(local),
701 window("2026-02-28T00:00:00Z", "2026-03-01T00:00:00Z")
702 );
703 assert_eq!(
704 Period::Day.window_at(local),
705 Period::Day.window_at(local.to_offset(UtcOffset::UTC))
706 );
707 }
708
709 #[test]
710 fn anchored_periods_clamp_the_day_to_the_target_month() {
711 let cases = [
715 (
717 "2026-01-31T09:30:00Z",
718 1,
719 "2026-02-28T10:00:00Z",
720 "2026-02-28T09:30:00Z",
721 "2026-03-31T09:30:00Z",
722 ),
723 (
725 "2026-01-31T09:30:00Z",
726 1,
727 "2026-04-05T00:00:00Z",
728 "2026-03-31T09:30:00Z",
729 "2026-04-30T09:30:00Z",
730 ),
731 (
733 "2026-01-30T09:00:00Z",
734 1,
735 "2026-02-10T00:00:00Z",
736 "2026-01-30T09:00:00Z",
737 "2026-02-28T09:00:00Z",
738 ),
739 (
740 "2026-01-30T09:00:00Z",
741 1,
742 "2026-03-10T00:00:00Z",
743 "2026-02-28T09:00:00Z",
744 "2026-03-30T09:00:00Z",
745 ),
746 (
749 "2024-01-29T00:00:00Z",
750 1,
751 "2024-02-05T00:00:00Z",
752 "2024-01-29T00:00:00Z",
753 "2024-02-29T00:00:00Z",
754 ),
755 (
756 "2025-01-29T00:00:00Z",
757 1,
758 "2025-02-28T00:00:00Z",
759 "2025-02-28T00:00:00Z",
760 "2025-03-29T00:00:00Z",
761 ),
762 (
764 "2026-01-31T00:00:00Z",
765 3,
766 "2026-02-15T00:00:00Z",
767 "2026-01-31T00:00:00Z",
768 "2026-04-30T00:00:00Z",
769 ),
770 (
771 "2026-01-31T00:00:00Z",
772 3,
773 "2026-05-01T00:00:00Z",
774 "2026-04-30T00:00:00Z",
775 "2026-07-31T00:00:00Z",
776 ),
777 ];
778 for (anchor, every, now, start, end) in cases {
779 let period = Period::Anchored {
780 anchor: at(anchor),
781 every: Months(every),
782 };
783 assert_eq!(
784 period.window_at(at(now)),
785 window(start, end),
786 "anchor {anchor}, every {every}, now {now}"
787 );
788 }
789
790 let period = Period::Anchored {
792 anchor: at("2026-03-15T08:00:00Z"),
793 every: Months(1),
794 };
795 assert_eq!(
796 period.window_at(at("2026-01-01T00:00:00Z")),
797 window("2025-12-15T08:00:00Z", "2026-01-15T08:00:00Z")
798 );
799 assert_eq!(period.index_at(at("2026-01-01T00:00:00Z")), -3);
800 }
801
802 #[test]
803 fn index_and_window_are_inverses() {
804 for period in [
805 Period::CalendarMonthUtc,
806 Period::Day,
807 Period::Anchored {
808 anchor: at("2026-01-31T09:30:00Z"),
809 every: Months(1),
810 },
811 ] {
812 let now = at("2026-08-31T09:30:00Z");
813 let index = period.index_at(now);
814 assert_eq!(period.window_for(index), period.window_at(now));
815 assert_eq!(
816 period.window_for(index - 1).end,
817 period.window_for(index).start,
818 "consecutive windows abut"
819 );
820 }
821 }
822
823 #[test]
831 fn a_period_index_outside_the_calendar_clamps_instead_of_panicking() {
832 let now = at("2026-01-15T00:00:00Z");
833 for index in [i64::MIN, i64::MIN + 1, -1_000_000, 1_000_000, i64::MAX] {
834 let month = Period::CalendarMonthUtc.window_for(index);
835 assert!(
836 month.start <= month.end,
837 "index {index}: {} .. {}",
838 month.start,
839 month.end
840 );
841 let day = Period::Day.window_for(index);
842 assert!(
843 day.start <= day.end,
844 "index {index}: {} .. {}",
845 day.start,
846 day.end
847 );
848 }
849 let current = Period::CalendarMonthUtc.index_at(now);
852 assert!(
853 Period::CalendarMonthUtc
854 .window_for(current - i64::from(u32::MAX) + 1)
855 .start
856 <= now
857 );
858 }
859
860 #[test]
865 fn an_anchor_spanning_more_months_than_a_calendar_holds_does_not_panic() {
866 for every in [Months(u32::MAX), Months(200_000), Months(120_000)] {
867 let period = Period::Anchored {
868 anchor: at("2026-01-31T09:30:00Z"),
869 every,
870 };
871 for index in [0, 1, -1, i64::MAX, i64::MIN] {
872 let window = period.window_for(index);
873 assert!(window.start <= window.end, "every {every:?}, index {index}");
874 }
875 let index = period.index_at(at("2026-01-15T00:00:00Z"));
876 assert!(
877 period.window_for(index).start <= at("2026-01-15T00:00:00Z"),
878 "every {every:?}"
879 );
880 }
881 }
882}