Limit&Quota Saving as File + API

src/quota_state.rs

@@ -0,0 +1,114 @@
+use std::collections::BTreeMap;
+use std::path::Path;
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use serde::{Deserialize, Serialize};
+use tokio::io::AsyncWriteExt;
+use tracing::{info, warn};
+
+use crate::stats::{Stats, UserQuotaSnapshot};
+
+#[derive(Debug, Default, Serialize, Deserialize)]
+pub(crate) struct QuotaStateFile {
+    pub(crate) last_reset_epoch_secs: u64,
+    pub(crate) users: BTreeMap<String, QuotaUserState>,
+}
+
+#[derive(Debug, Default, Serialize, Deserialize)]
+pub(crate) struct QuotaUserState {
+    pub(crate) used_bytes: u64,
+    pub(crate) last_reset_epoch_secs: u64,
+}
+
+fn now_epoch_secs() -> u64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_secs()
+}
+
+pub(crate) async fn load_quota_state(path: &Path, stats: &Stats) {
+    let bytes = match tokio::fs::read(path).await {
+        Ok(bytes) => bytes,
+        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return,
+        Err(error) => {
+            warn!(
+                error = %error,
+                path = %path.display(),
+                "Failed to read quota state file"
+            );
+            return;
+        }
+    };
+
+    let state = match serde_json::from_slice::<QuotaStateFile>(&bytes) {
+        Ok(state) => state,
+        Err(error) => {
+            warn!(
+                error = %error,
+                path = %path.display(),
+                "Failed to parse quota state file"
+            );
+            return;
+        }
+    };
+
+    let loaded_users = state.users.len();
+    for (user, quota) in state.users {
+        stats.load_user_quota_state(&user, quota.used_bytes, quota.last_reset_epoch_secs);
+    }
+    info!(
+        path = %path.display(),
+        loaded_users,
+        "Loaded per-user quota state"
+    );
+}
+
+pub(crate) async fn save_quota_state(path: &Path, stats: &Stats) -> std::io::Result<()> {
+    let mut users = BTreeMap::new();
+    let mut last_reset_epoch_secs = 0;
+    for (user, quota) in stats.user_quota_snapshot() {
+        last_reset_epoch_secs = last_reset_epoch_secs.max(quota.last_reset_epoch_secs);
+        users.insert(user, quota_user_state(quota));
+    }
+
+    let state = QuotaStateFile {
+        last_reset_epoch_secs,
+        users,
+    };
+    write_state_file(path, &state).await
+}
+
+pub(crate) async fn reset_user_quota(
+    path: &Path,
+    stats: &Stats,
+    user: &str,
+) -> std::io::Result<UserQuotaSnapshot> {
+    let snapshot = stats.reset_user_quota(user);
+    save_quota_state(path, stats).await?;
+    Ok(snapshot)
+}
+
+async fn write_state_file(path: &Path, state: &QuotaStateFile) -> std::io::Result<()> {
+    if let Some(parent) = path.parent()
+        && !parent.as_os_str().is_empty()
+    {
+        tokio::fs::create_dir_all(parent).await?;
+    }
+
+    let tmp_path = path.with_extension(format!("tmp.{}", now_epoch_secs()));
+    let payload = serde_json::to_vec_pretty(state)?;
+    let mut file = tokio::fs::File::create(&tmp_path).await?;
+    file.write_all(&payload).await?;
+    file.write_all(b"\n").await?;
+    file.sync_all().await?;
+    drop(file);
+    tokio::fs::rename(&tmp_path, path).await
+}
+
+fn quota_user_state(quota: UserQuotaSnapshot) -> QuotaUserState {
+    QuotaUserState {
+        used_bytes: quota.used_bytes,
+        last_reset_epoch_secs: quota.last_reset_epoch_secs,
+    }
+}