refactor: 将限流功能集成至Raft缓存管理器

heqingpan · heqingpan · commit 20db4961a805 · 2026-01-14T08:58:05.000+08:00
- 在CacheManagerRaftReq和CacheManagerRaftResult中添加Limit和Limiter变体
- 在DirectCacheManager中实现handle_limit方法处理限流逻辑
- 更新Console登录API和OpenAPI登录使用新的Raft缓存限流
- 统一登录限流和限流计数清除的缓存操作
- 更新AGENTS.md构建命令
diff --git a/AGENTS.md b/AGENTS.md
@@ -1,7 +1,7 @@
 # AGENTS.md - r-nacos 开发指南
 
 ## 构建/测试命令
-- 构建：`cargo build --release`
+- 构建：`cargo build`
 - 测试：`python3 integration_tests/scripts/run_test.py` (运行完整集成测试)
 - 单元测试：`cargo test` (运行所有单元测试)
 - 单个测试：`cargo test test_name` (运行指定测试)
@@ -42,4 +42,4 @@
 ### 部署相关
 - 环境变量前缀：`RNACOS_`
 - 默认端口：8848 (HTTP), 9848 (gRPC), 10848 (Console)
-- 支持单机和集群部署模式
+- 支持单机和集群部署模式
diff --git a/src/cache/actor_model.rs b/src/cache/actor_model.rs
@@ -1,7 +1,10 @@
 use crate::cache::model::{CacheKey, CacheValue};
 use crate::now_second_i32;
+use crate::raft::cache::CacheLimiterReq;
 use actix::Message;
+use ratelimiter_rs::RateLimiter;
 use serde::{Deserialize, Serialize};
+use std::sync::Arc;
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct CacheSetParam {
@@ -60,6 +63,7 @@ pub enum CacheManagerRaftReq {
     Get(CacheKey),
     Exists(CacheKey),
     Ttl(CacheKey),
+    Limit(CacheLimiterReq),
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -70,6 +74,7 @@ pub enum CacheManagerRaftResult {
     Value(CacheValue),
     Exists(bool),
     Ttl(i32),
+    Limiter(bool),
 }
 
 pub type CacheManagerResult = CacheManagerRaftResult;
diff --git a/src/cache/core.rs b/src/cache/core.rs
@@ -1,18 +1,24 @@
 use crate::cache::actor_model::{
-    CacheManagerLocalReq, CacheManagerRaftReq, CacheManagerRaftResult, CacheSetParam,
+    CacheManagerLocalReq, CacheManagerRaftReq, CacheManagerRaftResult, CacheManagerResult,
+    CacheSetParam,
 };
 use crate::cache::model::{CacheKey, CacheValue};
 use crate::common::constant::DIRECT_CACHE_TABLE_NAME;
+use crate::common::datetime_utils::now_millis_i64;
 use crate::common::datetime_utils::now_second_i32;
+use crate::common::limiter_utils::LimiterData;
 use crate::common::pb::data_object::DirectCacheItemDo;
+use crate::raft::cache::CacheLimiterReq;
 use crate::raft::filestore::model::SnapshotRecordDto;
 use crate::raft::filestore::raftapply::{RaftApplyDataRequest, RaftApplyDataResponse};
 use crate::raft::filestore::raftsnapshot::{SnapshotWriterActor, SnapshotWriterRequest};
 use actix::prelude::*;
 use bean_factory::{bean, Inject};
-use inner_mem_cache::{MemCache, TimeoutSet};
+use inner_mem_cache::TimeoutSet;
 use quick_protobuf::{BytesReader, Writer};
+use ratelimiter_rs::RateLimiter;
 use std::collections::HashMap;
+use std::convert::TryInto;
 use std::sync::Arc;
 
 pub struct CacheItem {
@@ -211,6 +217,43 @@ impl DirectCacheManager {
         }
     }
 
+    fn handle_limit(&mut self, limit_req: CacheLimiterReq) -> anyhow::Result<CacheManagerResult> {
+        let (rate_to_ms_conversion, key, limit) = match limit_req {
+            CacheLimiterReq::Second { key, limit } => (1000, key, limit),
+            CacheLimiterReq::Minutes { key, limit } => (60 * 1000, key, limit),
+            CacheLimiterReq::Hour { key, limit } => (60 * 60 * 1000, key, limit),
+            CacheLimiterReq::Day { key, limit } => (24 * 60 * 60 * 1000, key, limit),
+            CacheLimiterReq::OtherMills {
+                key,
+                limit,
+                rate_to_ms_conversion,
+            } => (rate_to_ms_conversion, key, limit),
+        };
+        let key = CacheKey::new(crate::cache::model::CacheType::String, key);
+        let now = now_millis_i64();
+        let mut limiter = if let Some(v) = self.get_valid_value(&key) {
+            if let Some(s) = v.try_to_string() {
+                let limiter_data: LimiterData = s.as_str().try_into()?;
+                limiter_data.to_rate_limiter()
+            } else {
+                RateLimiter::load(rate_to_ms_conversion, 0, now)
+            }
+        } else {
+            RateLimiter::load(rate_to_ms_conversion, 0, now)
+        };
+        let r = limiter.acquire(limit, limit as i64);
+        if r {
+            let limiter_data: LimiterData = limiter.into();
+            let cache_value = CacheValue::String(Arc::new(limiter_data.to_string()));
+            self.do_set(
+                key.clone(),
+                cache_value,
+                ((now + rate_to_ms_conversion as i64) / 1000) as i32,
+            );
+        }
+        Ok(CacheManagerRaftResult::Limiter(r))
+    }
+
     fn build_snapshot(&self, writer: Addr<SnapshotWriterActor>) -> anyhow::Result<()> {
         let now = now_second_i32();
         for (key, v) in self.cache.iter() {
@@ -295,6 +338,7 @@ impl Handler<CacheManagerRaftReq> for DirectCacheManager {
             CacheManagerRaftReq::Ttl(key) => Ok(self.get_ttl(&key)),
             CacheManagerRaftReq::Incr(key, expire) => Ok(self.incr(key, expire)),
             CacheManagerRaftReq::Decr(key, expire) => Ok(self.decr(key, expire)),
+            CacheManagerRaftReq::Limit(limit_req) => self.handle_limit(limit_req),
         }
     }
 }
diff --git a/src/console/login_api.rs b/src/console/login_api.rs
@@ -7,7 +7,7 @@ use crate::ldap::model::actor_model::{LdapMsgReq, LdapMsgResult};
 use crate::ldap::model::LdapUserParam;
 use crate::oauth2::model::actor_model::{OAuth2MsgReq, OAuth2MsgResult};
 use crate::oauth2::model::OAuth2UserParam;
-use crate::raft::store::ClientRequest;
+use crate::raft::store::{ClientRequest, ClientResponse};
 use crate::{
     common::{
         appdata::AppShareData,
@@ -255,8 +255,16 @@ async fn apply_session(
         .request(ClientRequest::CacheReq { req: cache_req })
         .await?;
     //登录成功后清除登陆限流计数
-    let clear_limit_req = CacheManagerReq::Remove(CacheKey::new(CacheType::String, limit_key));
-    app.cache_manager.do_send(clear_limit_req);
+    let clear_limit_req = crate::cache::actor_model::CacheManagerRaftReq::Remove(CacheKey::new(
+        CacheType::String,
+        limit_key,
+    ));
+    app.raft_request_route
+        .request(ClientRequest::CacheReq {
+            req: clear_limit_req,
+        })
+        .await
+        .ok();
     let login_token = LoginToken {
         token: token.to_string(),
     };
@@ -315,9 +323,14 @@ async fn login_limit(
         key: limit_key.clone(),
         limit: app.sys_config.console_login_one_hour_limit as i32,
     };
+    let cache_req = crate::cache::actor_model::CacheManagerRaftReq::Limit(limit_req);
     //登录前先判断是否登陆准入
-    if let Ok(CacheManagerResult::Limiter(acquire_result)) =
-        app.raft_cache_route.request_limiter(limit_req).await
+    if let Ok(ClientResponse::CacheResp {
+        resp: CacheManagerRaftResult::Limiter(acquire_result),
+    }) = app
+        .raft_request_route
+        .request(ClientRequest::CacheReq { req: cache_req })
+        .await
     {
         if !acquire_result {
             return Some(HttpResponse::Ok().json(ApiResult::<()>::error(
diff --git a/src/openapi/auth.rs b/src/openapi/auth.rs
@@ -1,11 +1,11 @@
-use crate::cache::actor_model::CacheSetParam;
+use crate::cache::actor_model::{CacheManagerRaftResult, CacheSetParam};
 use crate::common::appdata::AppShareData;
 use crate::common::model::TokenSession;
 use crate::common::option_utils::OptionUtils;
 use crate::merge_web_param_with_result;
 use crate::raft::cache::model::{CacheKey, CacheType, CacheValue};
 use crate::raft::cache::{CacheLimiterReq, CacheManagerReq, CacheManagerResult};
-use crate::raft::store::ClientRequest;
+use crate::raft::store::{ClientRequest, ClientResponse};
 use crate::user::{UserManagerReq, UserManagerResult};
 use actix_web::{web, HttpResponse, Responder};
 use serde::{Deserialize, Serialize};
@@ -70,9 +70,14 @@ async fn do_login(
         key: limit_key.clone(),
         limit: app.sys_config.openapi_login_one_minute_limit as i32,
     };
+    let cache_req = crate::cache::actor_model::CacheManagerRaftReq::Limit(limit_req);
     //登录前先判断是否登陆准入
-    if let Ok(CacheManagerResult::Limiter(acquire_result)) =
-        app.raft_cache_route.request_limiter(limit_req).await
+    if let ClientResponse::CacheResp {
+        resp: CacheManagerRaftResult::Limiter(acquire_result),
+    } = app
+        .raft_request_route
+        .request(ClientRequest::CacheReq { req: cache_req })
+        .await?
     {
         if !acquire_result {
             return Err(anyhow::anyhow!(
@@ -118,9 +123,15 @@ async fn do_login(
                 .request(ClientRequest::CacheReq { req: cache_req })
                 .await?;
             //登录成功后清除登陆限流计数
-            let clear_limit_req =
-                CacheManagerReq::Remove(CacheKey::new(CacheType::String, limit_key));
-            app.cache_manager.do_send(clear_limit_req);
+            let clear_limit_req = crate::cache::actor_model::CacheManagerRaftReq::Remove(
+                CacheKey::new(CacheType::String, limit_key),
+            );
+            app.raft_request_route
+                .request(ClientRequest::CacheReq {
+                    req: clear_limit_req,
+                })
+                .await
+                .ok();
             let login_result = LoginResult {
                 access_token: Some(token),
                 token_ttl: app.sys_config.openapi_login_timeout as i64,
