feat: implement bucket quota system (#1461)

Signed-off-by: yxrxy <1532529704@qq.com> Co-authored-by: loverustfs <hello@rustfs.com>
2026-08-28 07:57:01 +00:00 · 2026-01-12 11:42:07 +08:00
parent 78b13f3ff2
commit 29d86036b1
17 changed files with 1964 additions and 42 deletions
@@ -355,7 +355,7 @@ impl BucketMetadata {
            self.tagging_config = Some(deserialize::<Tagging>(&self.tagging_config_xml)?);
        }
        if !self.quota_config_json.is_empty() {
-            self.quota_config = Some(BucketQuota::unmarshal(&self.quota_config_json)?);
+            self.quota_config = Some(serde_json::from_slice(&self.quota_config_json)?);
        }
        if !self.replication_config_xml.is_empty() {
            self.replication_config = Some(deserialize::<ReplicationConfiguration>(&self.replication_config_xml)?);
@@ -487,7 +487,8 @@ mod test {
        bm.tagging_config_updated_at = OffsetDateTime::now_utc();

        // Add quota configuration
-        let quota_json = r#"{"quota":1073741824,"quotaType":"hard"}"#; // 1GB quota
+        let quota_json =
+            r#"{"quota":1073741824,"quota_type":"Hard","created_at":"2024-01-01T00:00:00Z","updated_at":"2024-01-01T00:00:00Z"}"#; // 1GB quota
        bm.quota_config_json = quota_json.as_bytes().to_vec();
        bm.quota_config_updated_at = OffsetDateTime::now_utc();

@@ -0,0 +1,195 @@
+// Copyright 2024 RustFS Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::{BucketQuota, QuotaCheckResult, QuotaError, QuotaOperation};
+use crate::bucket::metadata_sys::{BucketMetadataSys, update};
+use crate::data_usage::get_bucket_usage_memory;
+use rustfs_common::metrics::Metric;
+use rustfs_config::QUOTA_CONFIG_FILE;
+use std::sync::Arc;
+use std::time::Instant;
+use tokio::sync::RwLock;
+use tracing::{debug, warn};
+
+pub struct QuotaChecker {
+    metadata_sys: Arc<RwLock<BucketMetadataSys>>,
+}
+
+impl QuotaChecker {
+    pub fn new(metadata_sys: Arc<RwLock<BucketMetadataSys>>) -> Self {
+        Self { metadata_sys }
+    }
+
+    pub async fn check_quota(
+        &self,
+        bucket: &str,
+        operation: QuotaOperation,
+        operation_size: u64,
+    ) -> Result<QuotaCheckResult, QuotaError> {
+        let start_time = Instant::now();
+        let quota_config = self.get_quota_config(bucket).await?;
+
+        // If no quota limit is set, allow operation
+        let quota_limit = match quota_config.quota {
+            None => {
+                let current_usage = self.get_real_time_usage(bucket).await?;
+                return Ok(QuotaCheckResult {
+                    allowed: true,
+                    current_usage,
+                    quota_limit: None,
+                    operation_size,
+                    remaining: None,
+                });
+            }
+            Some(q) => q,
+        };
+
+        let current_usage = self.get_real_time_usage(bucket).await?;
+
+        let expected_usage = match operation {
+            QuotaOperation::PutObject | QuotaOperation::CopyObject => current_usage + operation_size,
+            QuotaOperation::DeleteObject => current_usage.saturating_sub(operation_size),
+        };
+
+        let allowed = match operation {
+            QuotaOperation::PutObject | QuotaOperation::CopyObject => {
+                quota_config.check_operation_allowed(current_usage, operation_size)
+            }
+            QuotaOperation::DeleteObject => true,
+        };
+
+        let remaining = if quota_limit >= expected_usage {
+            Some(quota_limit - expected_usage)
+        } else {
+            Some(0)
+        };
+
+        if !allowed {
+            warn!(
+                "Quota exceeded for bucket: {}, current: {}, limit: {}, attempted: {}",
+                bucket, current_usage, quota_limit, operation_size
+            );
+        }
+
+        let result = QuotaCheckResult {
+            allowed,
+            current_usage,
+            quota_limit: Some(quota_limit),
+            operation_size,
+            remaining,
+        };
+
+        let duration = start_time.elapsed();
+        rustfs_common::metrics::Metrics::inc_time(Metric::QuotaCheck, duration).await;
+        if !allowed {
+            rustfs_common::metrics::Metrics::inc_time(Metric::QuotaViolation, duration).await;
+        }
+
+        Ok(result)
+    }
+
+    pub async fn get_quota_config(&self, bucket: &str) -> Result<BucketQuota, QuotaError> {
+        let meta = self
+            .metadata_sys
+            .read()
+            .await
+            .get(bucket)
+            .await
+            .map_err(QuotaError::StorageError)?;
+
+        if meta.quota_config_json.is_empty() {
+            debug!("No quota config found for bucket: {}, using default", bucket);
+            return Ok(BucketQuota::new(None));
+        }
+
+        let quota: BucketQuota = serde_json::from_slice(&meta.quota_config_json).map_err(|e| QuotaError::InvalidConfig {
+            reason: format!("Failed to parse quota config: {}", e),
+        })?;
+
+        Ok(quota)
+    }
+
+    pub async fn set_quota_config(&mut self, bucket: &str, quota: BucketQuota) -> Result<(), QuotaError> {
+        let json_data = serde_json::to_vec(&quota).map_err(|e| QuotaError::InvalidConfig {
+            reason: format!("Failed to serialize quota config: {}", e),
+        })?;
+        let start_time = Instant::now();
+
+        update(bucket, QUOTA_CONFIG_FILE, json_data)
+            .await
+            .map_err(QuotaError::StorageError)?;
+
+        rustfs_common::metrics::Metrics::inc_time(Metric::QuotaSync, start_time.elapsed()).await;
+        Ok(())
+    }
+
+    pub async fn get_quota_stats(&self, bucket: &str) -> Result<(BucketQuota, Option<u64>), QuotaError> {
+        // If bucket doesn't exist, return ConfigNotFound error
+        if !self.bucket_exists(bucket).await {
+            return Err(QuotaError::ConfigNotFound {
+                bucket: bucket.to_string(),
+            });
+        }
+
+        let quota = self.get_quota_config(bucket).await?;
+        let current_usage = self.get_real_time_usage(bucket).await.unwrap_or(0);
+
+        Ok((quota, Some(current_usage)))
+    }
+
+    pub async fn bucket_exists(&self, bucket: &str) -> bool {
+        self.metadata_sys.read().await.get(bucket).await.is_ok()
+    }
+
+    pub async fn get_real_time_usage(&self, bucket: &str) -> Result<u64, QuotaError> {
+        Ok(get_bucket_usage_memory(bucket).await.unwrap_or(0))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_quota_check_no_limit() {
+        let result = QuotaCheckResult {
+            allowed: true,
+            current_usage: 0,
+            quota_limit: None,
+            operation_size: 1024,
+            remaining: None,
+        };
+
+        assert!(result.allowed);
+        assert_eq!(result.quota_limit, None);
+    }
+
+    #[tokio::test]
+    async fn test_quota_check_within_limit() {
+        let quota = BucketQuota::new(Some(2048)); // 2KB
+
+        // Current usage 512, trying to add 1024
+        let allowed = quota.check_operation_allowed(512, 1024);
+        assert!(allowed);
+    }
+
+    #[tokio::test]
+    async fn test_quota_check_exceeds_limit() {
+        let quota = BucketQuota::new(Some(1024)); // 1KB
+
+        // Current usage 512, trying to add 1024
+        let allowed = quota.check_operation_allowed(512, 1024);
+        assert!(!allowed);
+    }
+}
@@ -12,36 +12,37 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.

+pub mod checker;
+
 use crate::error::Result;
 use rmp_serde::Serializer as rmpSerializer;
+use rustfs_config::{
+    QUOTA_API_PATH, QUOTA_EXCEEDED_ERROR_CODE, QUOTA_INTERNAL_ERROR_CODE, QUOTA_INVALID_CONFIG_ERROR_CODE,
+    QUOTA_NOT_FOUND_ERROR_CODE,
+};
 use serde::{Deserialize, Serialize};
+use thiserror::Error;
+use time::OffsetDateTime;

-// Define the QuotaType enum
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
 pub enum QuotaType {
+    /// Hard quota: reject immediately when exceeded
+    #[default]
    Hard,
 }

-// Define the BucketQuota structure
-#[derive(Debug, Deserialize, Serialize, Default, Clone)]
+#[derive(Debug, Deserialize, Serialize, Default, Clone, PartialEq)]
 pub struct BucketQuota {
-    quota: Option<u64>, // Use Option to represent optional fields
-
-    size: u64,
-
-    rate: u64,
-
-    requests: u64,
-
-    quota_type: Option<QuotaType>,
+    pub quota: Option<u64>,
+    pub quota_type: QuotaType,
+    /// Timestamp when this quota configuration was set (for audit purposes)
+    pub created_at: Option<OffsetDateTime>,
 }

 impl BucketQuota {
    pub fn marshal_msg(&self) -> Result<Vec<u8>> {
        let mut buf = Vec::new();
-
        self.serialize(&mut rmpSerializer::new(&mut buf).with_struct_map())?;
-
        Ok(buf)
    }

@@ -49,4 +50,107 @@ impl BucketQuota {
        let t: BucketQuota = rmp_serde::from_slice(buf)?;
        Ok(t)
    }
+
+    pub fn new(quota: Option<u64>) -> Self {
+        let now = OffsetDateTime::now_utc();
+        Self {
+            quota,
+            quota_type: QuotaType::Hard,
+            created_at: Some(now),
+        }
+    }
+
+    pub fn get_quota_limit(&self) -> Option<u64> {
+        self.quota
+    }
+
+    pub fn check_operation_allowed(&self, current_usage: u64, operation_size: u64) -> bool {
+        if let Some(quota_limit) = self.quota {
+            current_usage.saturating_add(operation_size) <= quota_limit
+        } else {
+            true // No quota limit
+        }
+    }
+
+    pub fn get_remaining_quota(&self, current_usage: u64) -> Option<u64> {
+        self.quota.map(|limit| limit.saturating_sub(current_usage))
+    }
+}
+
+#[derive(Debug)]
+pub struct QuotaCheckResult {
+    pub allowed: bool,
+    pub current_usage: u64,
+    /// quota_limit: None means unlimited
+    pub quota_limit: Option<u64>,
+    pub operation_size: u64,
+    pub remaining: Option<u64>,
+}
+
+#[derive(Debug)]
+pub enum QuotaOperation {
+    PutObject,
+    CopyObject,
+    DeleteObject,
+}
+
+#[derive(Debug, Error)]
+pub enum QuotaError {
+    #[error("Bucket quota exceeded: current={current}, limit={limit}, operation={operation}")]
+    QuotaExceeded { current: u64, limit: u64, operation: u64 },
+    #[error("Quota configuration not found for bucket: {bucket}")]
+    ConfigNotFound { bucket: String },
+    #[error("Invalid quota configuration: {reason}")]
+    InvalidConfig { reason: String },
+    #[error("Storage error: {0}")]
+    StorageError(#[from] crate::error::StorageError),
+}
+
+#[derive(Debug, Serialize)]
+pub struct QuotaErrorResponse {
+    #[serde(rename = "Code")]
+    pub code: String,
+    #[serde(rename = "Message")]
+    pub message: String,
+    #[serde(rename = "Resource")]
+    pub resource: String,
+    #[serde(rename = "RequestId")]
+    pub request_id: String,
+    #[serde(rename = "HostId")]
+    pub host_id: String,
+}
+
+impl QuotaErrorResponse {
+    pub fn new(quota_error: &QuotaError, request_id: &str, host_id: &str) -> Self {
+        match quota_error {
+            QuotaError::QuotaExceeded { .. } => Self {
+                code: QUOTA_EXCEEDED_ERROR_CODE.to_string(),
+                message: quota_error.to_string(),
+                resource: QUOTA_API_PATH.to_string(),
+                request_id: request_id.to_string(),
+                host_id: host_id.to_string(),
+            },
+            QuotaError::ConfigNotFound { .. } => Self {
+                code: QUOTA_NOT_FOUND_ERROR_CODE.to_string(),
+                message: quota_error.to_string(),
+                resource: QUOTA_API_PATH.to_string(),
+                request_id: request_id.to_string(),
+                host_id: host_id.to_string(),
+            },
+            QuotaError::InvalidConfig { .. } => Self {
+                code: QUOTA_INVALID_CONFIG_ERROR_CODE.to_string(),
+                message: quota_error.to_string(),
+                resource: QUOTA_API_PATH.to_string(),
+                request_id: request_id.to_string(),
+                host_id: host_id.to_string(),
+            },
+            QuotaError::StorageError(_) => Self {
+                code: QUOTA_INTERNAL_ERROR_CODE.to_string(),
+                message: quota_error.to_string(),
+                resource: QUOTA_API_PATH.to_string(),
+                request_id: request_id.to_string(),
+                host_id: host_id.to_string(),
+            },
+        }
+    }
 }