use crate::clock::Clock;
use crate::GatewardenError;
use chrono::{DateTime, Datelike, Utc};
use serde::{Deserialize, Serialize};
use std::fs;
use std::path::PathBuf;
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct UsageStats {
pub daily_count: u64,
pub monthly_count: u64,
pub daily_date: Option<String>,
pub monthly_period: Option<String>,
pub lifetime_count: u64,
}
impl UsageStats {
pub fn new() -> Self {
Self::default()
}
pub fn increment(&mut self, clock: &dyn Clock) {
let now = clock.now_utc();
let today = format_date(&now);
let this_month = format_month(&now);
if self.daily_date.as_ref() != Some(&today) {
self.daily_count = 0;
self.daily_date = Some(today);
}
if self.monthly_period.as_ref() != Some(&this_month) {
self.monthly_count = 0;
self.monthly_period = Some(this_month);
}
self.daily_count += 1;
self.monthly_count += 1;
self.lifetime_count += 1;
}
pub fn get_daily_count(&self, clock: &dyn Clock) -> u64 {
let now = clock.now_utc();
let today = format_date(&now);
if self.daily_date.as_ref() == Some(&today) {
self.daily_count
} else {
0
}
}
pub fn get_monthly_count(&self, clock: &dyn Clock) -> u64 {
let now = clock.now_utc();
let this_month = format_month(&now);
if self.monthly_period.as_ref() == Some(&this_month) {
self.monthly_count
} else {
0
}
}
}
fn format_date(dt: &DateTime<Utc>) -> String {
format!("{:04}-{:02}-{:02}", dt.year(), dt.month(), dt.day())
}
fn format_month(dt: &DateTime<Utc>) -> String {
format!("{:04}-{:02}", dt.year(), dt.month())
}
pub struct UsageMeter {
path: PathBuf,
stats: UsageStats,
}
impl UsageMeter {
pub fn new(path: PathBuf) -> Result<Self, GatewardenError> {
let stats = if path.exists() {
let json = fs::read_to_string(&path)
.map_err(|e| GatewardenError::MeterIO(format!("Failed to read meter: {}", e)))?;
serde_json::from_str(&json)
.map_err(|e| GatewardenError::MeterIO(format!("Failed to parse meter: {}", e)))?
} else {
UsageStats::new()
};
Ok(Self { path, stats })
}
pub fn with_namespace(namespace: &str) -> Result<Self, GatewardenError> {
let base_dir = dirs::data_dir()
.ok_or_else(|| GatewardenError::MeterIO("Could not find data directory".to_string()))?;
let dir = base_dir.join(namespace);
fs::create_dir_all(&dir)
.map_err(|e| GatewardenError::MeterIO(format!("Failed to create dir: {}", e)))?;
let path = dir.join("usage.json");
Self::new(path)
}
pub fn increment(&mut self, clock: &dyn Clock) -> Result<(), GatewardenError> {
self.stats.increment(clock);
self.save()
}
pub fn daily_count(&self, clock: &dyn Clock) -> u64 {
self.stats.get_daily_count(clock)
}
pub fn monthly_count(&self, clock: &dyn Clock) -> u64 {
self.stats.get_monthly_count(clock)
}
pub fn lifetime_count(&self) -> u64 {
self.stats.lifetime_count
}
pub fn stats(&self) -> &UsageStats {
&self.stats
}
fn save(&self) -> Result<(), GatewardenError> {
if let Some(parent) = self.path.parent() {
fs::create_dir_all(parent)
.map_err(|e| GatewardenError::MeterIO(format!("Failed to create dir: {}", e)))?;
}
let json = serde_json::to_string_pretty(&self.stats)
.map_err(|e| GatewardenError::MeterIO(format!("Failed to serialize: {}", e)))?;
let temp_path = self.path.with_extension("tmp");
fs::write(&temp_path, &json)
.map_err(|e| GatewardenError::MeterIO(format!("Failed to write temp: {}", e)))?;
fs::rename(&temp_path, &self.path)
.map_err(|e| GatewardenError::MeterIO(format!("Failed to rename: {}", e)))?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::clock::MockClock;
use chrono::TimeZone;
use tempfile::TempDir;
#[test]
fn test_usage_stats_increment() {
let clock = MockClock::new(Utc.with_ymd_and_hms(2025, 1, 15, 12, 0, 0).unwrap());
let mut stats = UsageStats::new();
stats.increment(&clock);
assert_eq!(stats.daily_count, 1);
assert_eq!(stats.monthly_count, 1);
assert_eq!(stats.lifetime_count, 1);
stats.increment(&clock);
assert_eq!(stats.daily_count, 2);
assert_eq!(stats.monthly_count, 2);
assert_eq!(stats.lifetime_count, 2);
}
#[test]
fn test_usage_stats_daily_rollover() {
let clock1 = MockClock::new(Utc.with_ymd_and_hms(2025, 1, 15, 12, 0, 0).unwrap());
let mut stats = UsageStats::new();
stats.increment(&clock1);
stats.increment(&clock1);
assert_eq!(stats.daily_count, 2);
let clock2 = MockClock::new(Utc.with_ymd_and_hms(2025, 1, 16, 12, 0, 0).unwrap());
stats.increment(&clock2);
assert_eq!(stats.daily_count, 1);
assert_eq!(stats.monthly_count, 3); assert_eq!(stats.lifetime_count, 3);
}
#[test]
fn test_usage_stats_monthly_rollover() {
let clock1 = MockClock::new(Utc.with_ymd_and_hms(2025, 1, 31, 23, 59, 0).unwrap());
let mut stats = UsageStats::new();
stats.increment(&clock1);
assert_eq!(stats.daily_count, 1);
assert_eq!(stats.monthly_count, 1);
let clock2 = MockClock::new(Utc.with_ymd_and_hms(2025, 2, 1, 0, 0, 0).unwrap());
stats.increment(&clock2);
assert_eq!(stats.daily_count, 1); assert_eq!(stats.monthly_count, 1); assert_eq!(stats.lifetime_count, 2);
}
#[test]
fn test_usage_stats_year_rollover() {
let clock1 = MockClock::new(Utc.with_ymd_and_hms(2024, 12, 31, 23, 59, 0).unwrap());
let mut stats = UsageStats::new();
stats.increment(&clock1);
let clock2 = MockClock::new(Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap());
stats.increment(&clock2);
assert_eq!(stats.daily_count, 1);
assert_eq!(stats.monthly_count, 1);
assert_eq!(stats.lifetime_count, 2);
}
#[test]
fn test_get_counts_with_rollover() {
let clock1 = MockClock::new(Utc.with_ymd_and_hms(2025, 1, 15, 12, 0, 0).unwrap());
let mut stats = UsageStats::new();
stats.increment(&clock1);
stats.increment(&clock1);
assert_eq!(stats.get_daily_count(&clock1), 2);
assert_eq!(stats.get_monthly_count(&clock1), 2);
let clock2 = MockClock::new(Utc.with_ymd_and_hms(2025, 1, 16, 12, 0, 0).unwrap());
assert_eq!(stats.get_daily_count(&clock2), 0);
assert_eq!(stats.get_monthly_count(&clock2), 2);
let clock3 = MockClock::new(Utc.with_ymd_and_hms(2025, 2, 1, 12, 0, 0).unwrap());
assert_eq!(stats.get_daily_count(&clock3), 0);
assert_eq!(stats.get_monthly_count(&clock3), 0);
}
#[test]
fn test_usage_meter_persistence() {
let temp_dir = TempDir::new().unwrap();
let path = temp_dir.path().join("usage.json");
let clock = MockClock::new(Utc.with_ymd_and_hms(2025, 1, 15, 12, 0, 0).unwrap());
{
let mut meter = UsageMeter::new(path.clone()).unwrap();
meter.increment(&clock).unwrap();
meter.increment(&clock).unwrap();
}
{
let meter = UsageMeter::new(path).unwrap();
assert_eq!(meter.daily_count(&clock), 2);
assert_eq!(meter.monthly_count(&clock), 2);
assert_eq!(meter.lifetime_count(), 2);
}
}
#[test]
fn test_usage_meter_daily_rollover() {
let temp_dir = TempDir::new().unwrap();
let path = temp_dir.path().join("usage.json");
let clock1 = MockClock::new(Utc.with_ymd_and_hms(2025, 1, 15, 12, 0, 0).unwrap());
{
let mut meter = UsageMeter::new(path.clone()).unwrap();
meter.increment(&clock1).unwrap();
meter.increment(&clock1).unwrap();
}
let clock2 = MockClock::new(Utc.with_ymd_and_hms(2025, 1, 16, 12, 0, 0).unwrap());
{
let mut meter = UsageMeter::new(path).unwrap();
meter.increment(&clock2).unwrap();
assert_eq!(meter.daily_count(&clock2), 1);
assert_eq!(meter.monthly_count(&clock2), 3);
assert_eq!(meter.lifetime_count(), 3);
}
}
#[test]
fn test_usage_meter_monthly_rollover() {
let temp_dir = TempDir::new().unwrap();
let path = temp_dir.path().join("usage.json");
let clock1 = MockClock::new(Utc.with_ymd_and_hms(2025, 1, 31, 23, 59, 0).unwrap());
{
let mut meter = UsageMeter::new(path.clone()).unwrap();
meter.increment(&clock1).unwrap();
meter.increment(&clock1).unwrap();
}
let clock2 = MockClock::new(Utc.with_ymd_and_hms(2025, 2, 1, 0, 0, 0).unwrap());
{
let mut meter = UsageMeter::new(path).unwrap();
meter.increment(&clock2).unwrap();
assert_eq!(meter.daily_count(&clock2), 1);
assert_eq!(meter.monthly_count(&clock2), 1);
assert_eq!(meter.lifetime_count(), 3);
}
}
}