Merge pull request #13 from cn-kali-team/main

multi-tenant

Author: cn-kali-team

asynq-server/Cargo.toml

@@ -29,7 +29,7 @@ tower = "0.5"
 tower-http = { version = "0.6", features = ["cors", "trace"] }
 
 # Async runtime
-tokio = { version = "1.0", features = ["rt", "macros", "signal", "rt-multi-thread", "sync"] }
+tokio = { version = "1.0", features = ["rt", "macros", "signal", "rt-multi-thread", "sync", "time"] }
 
 # WebSocket
 tokio-tungstenite = "0.28"

asynq/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "asynq"
-version = "0.1.5"
+version = "0.1.6"
 license.workspace = true
 edition.workspace = true
 authors.workspace = true
@@ -18,7 +18,7 @@ path = "src/lib.rs"
 [dependencies]
 # Redis 连接池和客户端（始终启用作为默认后端）
 # Redis connection pool and client (always enabled as the default backend)
-redis = { version = "1.0.2", features = ["tokio-comp", "aio", "connection-manager", "streams", "acl"] }
+redis = { version = "1.0.3", features = ["tokio-comp", "aio", "connection-manager", "streams", "acl"] }
 # 宏
 asynq-macros = { version = "0.1.0", optional = true }
 # 序列化和反序列化

asynq/examples/acl_example.rs

@@ -42,6 +42,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         Ok(exists) => {
           if exists {
             println!("   User '{}' already exists", node_config.username);
+            manager
+              .delete_tenant_user(&acl_config.node_config.username)
+              .await?;
+            match manager.create_tenant_user(&acl_config).await {
+              Ok(_) => println!("   ✅ Created user '{}'", node_config.username),
+              Err(e) => println!("   ❌ Failed to create user: {e}"),
+            }
           } else {
             println!("   User '{}' does not exist", node_config.username);
             // 创建租户用户（需要管理员权限）
@@ -78,15 +85,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     Err(e) => {
       println!("   ⚠️  Could not connect to Redis: {e}");
       println!("   Make sure Redis is running at: {redis_url}");
-      println!("\n   Showing ACL command that would be generated:");
-
-      // 即使没有连接，也可以显示配置信息
-      // Even without connection, we can show configuration info
-      println!("   ACL SETUSER {} on +@all -@dangerous +keys +info|memory +info|clients -select +select|{} >*** {:?}",
-        acl_config.node_config.username,
-        acl_config.node_config.db,
-        acl_config.node_config.asynq_key_pattern()
-      );
     }
   }
   Ok(())

asynq/src/acl/mod.rs

@@ -156,16 +156,29 @@ impl AclConfig {
 
   /// 获取默认的键模式列表
   /// Get default key pattern list
-  pub fn default_key_patterns() -> Vec<Rule> {
+  pub fn default_key_patterns(tenant: &str) -> Vec<Rule> {
     vec![
-      Rule::Pattern("asynq:queues".to_string()),
-      Rule::Pattern("asynq:servers:*".to_string()),
-      Rule::Pattern("asynq:servers".to_string()),
-      Rule::Pattern("asynq:workers".to_string()),
-      Rule::Pattern("asynq:workers:*".to_string()),
-      Rule::Pattern("asynq:schedulers".to_string()),
-      Rule::Pattern("asynq:schedulers:*".to_string()),
-      Rule::Other("&asynq:cancel".to_string()),
+      // Add default key patterns
+      Rule::Pattern(crate::base::keys::ALL_QUEUES.to_string()),
+      Rule::Pattern(format!(
+        "{}{{{}:*",
+        crate::base::keys::SERVERS_PREFIX,
+        tenant
+      )),
+      Rule::Pattern(crate::base::keys::ALL_SERVERS.to_string()),
+      Rule::Pattern(crate::base::keys::ALL_WORKERS.to_string()),
+      Rule::Pattern(format!(
+        "{}{{{}:*",
+        crate::base::keys::WORKERS_PREFIX,
+        tenant
+      )),
+      Rule::Pattern(crate::base::keys::ALL_SCHEDULERS.to_string()),
+      Rule::Pattern(format!(
+        "{}{{{}:*",
+        crate::base::keys::SCHEDULED_PREFIX,
+        tenant
+      )),
+      Rule::Channel(crate::base::keys::CANCEL_CHANNEL.to_string()),
     ]
   }
 }

asynq/src/acl/redis_acl.rs

@@ -73,37 +73,27 @@ impl RedisAclManager {
   /// 生成 ACL 规则列表
   /// Generate ACL rules list
   fn build_acl_rules(&self, config: &AclConfig) -> Vec<Rule> {
-    let mut rules = Vec::new();
-
-    // 基本权限: on, +@all, -@dangerous, +keys, +info|memory, +info|clients
-    // Basic permissions: on, +@all, -@dangerous, +keys, +info|memory, +info|clients
-    rules.push(Rule::On);
-    rules.push(Rule::AllCommands);
-    rules.push(Rule::RemoveCategory("dangerous".to_string()));
-    rules.push(Rule::AddCommand("keys".to_string()));
-    rules.push(Rule::RemoveCommand("info".to_string()));
-    // 数据库限制: -select, +select|<db>
-    // Database restrictions: -select, +select|<db>
-    rules.push(Rule::RemoveCommand("select".to_string()));
-    // 密码
-    // Password
-    rules.push(Rule::AddPass(config.node_config.password.clone()));
-
-    // 添加默认队列
-    // Add default queue pattern - uses hash tag {DEFAULT_QUEUE_NAME} for Redis cluster routing
-    rules.push(Rule::Pattern(format!("asynq:{{{}}}:*", DEFAULT_QUEUE_NAME)));
-
-    // 添加租户特定的键模式
-    // Add tenant-specific key patterns
-    // 从 "~asynq:{username:*" 格式提取模式部分（去掉前导的 ~）
-    rules.push(config.node_config.asynq_key_pattern());
-
-    // 添加默认的键模式
-    // Add default key patterns
-    for pattern in AclConfig::default_key_patterns() {
-      rules.push(pattern);
-    }
-
+    let mut rules = vec![
+      // Basic permissions: on, +@all, -@dangerous, +keys, -info
+      Rule::On,
+      Rule::ResetChannels,
+      Rule::AllCommands,
+      Rule::RemoveCategory("dangerous".to_string()),
+      Rule::RemoveCategory("admin".to_string()),
+      Rule::RemoveCommand("keys".to_string()),
+      Rule::RemoveCommand("info".to_string()),
+      // Database restrictions: -select
+      Rule::RemoveCommand("select".to_string()),
+      // Password
+      Rule::AddPass(config.node_config.password.clone()),
+      // Add default queue pattern - uses hashtag {DEFAULT_QUEUE_NAME} for Redis cluster routing
+      Rule::Pattern(format!("asynq:{{{}}}:*", DEFAULT_QUEUE_NAME)),
+      // Add tenant-specific key patterns
+      config.node_config.asynq_key_pattern(),
+    ];
+    rules.extend(AclConfig::default_key_patterns(
+      &config.node_config.username,
+    ));
     // 添加只写键
     // Add write-only keys
     for key in &config.write_only_keys {

asynq/src/backend/pgdb/broker.rs

@@ -942,7 +942,13 @@ impl Broker for PostgresBroker {
 
   /// 写入服务器状态。
   /// Write server state.
-  async fn write_server_state(&self, server_info: &ServerInfo, ttl: Duration) -> Result<()> {
+  async fn write_server_state(
+    &self,
+    server_info: &ServerInfo,
+    _worker: Vec<WorkerInfo>,
+    ttl: Duration,
+    _tenant: Option<&str>,
+  ) -> Result<()> {
     let server_key = format!(
       "{}:{}:{}",
       server_info.host, server_info.pid, server_info.server_id
@@ -992,7 +998,13 @@ impl Broker for PostgresBroker {
 
   /// 清除服务器状态。
   /// Clear server state.
-  async fn clear_server_state(&self, host: &str, pid: i32, server_id: &str) -> Result<()> {
+  async fn clear_server_state(
+    &self,
+    host: &str,
+    pid: i32,
+    server_id: &str,
+    _tenant: Option<&str>,
+  ) -> Result<()> {
     let server_key = format!("{host}:{pid}:{server_id}");
 
     // Clear all workers for this server

asynq/src/backend/rdb/broker.rs

@@ -5,7 +5,7 @@ use crate::base::keys;
 use crate::base::keys::TaskState;
 use crate::base::Broker;
 use crate::error::{Error, Result};
-use crate::proto::{ServerInfo, TaskMessage};
+use crate::proto::{ServerInfo, TaskMessage, WorkerInfo};
 use crate::task::{generate_task_id, Task, TaskInfo};
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
@@ -811,14 +811,23 @@ impl Broker for RedisBroker {
   }
 
   /// 写入服务器状态 - Go: WriteServerState
-  async fn write_server_state(&self, server_info: &ServerInfo, ttl: Duration) -> Result<()> {
+  async fn write_server_state(
+    &self,
+    server_info: &ServerInfo,
+    workers: Vec<WorkerInfo>,
+    ttl: Duration,
+    tenant: Option<&str>,
+  ) -> Result<()> {
     let mut conn = self.get_async_connection().await?;
 
     // 使用 ServerInfo 中的信息构建键
     // Build keys using information from ServerInfo
-    let server_key =
-      keys::server_info_key(&server_info.host, server_info.pid, &server_info.server_id);
-    let workers_key = keys::workers_key(&server_info.host, server_info.pid, &server_info.server_id);
+    let (server_key, workers_key) = keys::server_and_workers_keys(
+      tenant,
+      &server_info.host,
+      server_info.pid,
+      &server_info.server_id,
+    );
 
     // 计算过期时间戳
     // Calculate expiration timestamp
@@ -842,44 +851,44 @@ impl Broker for RedisBroker {
     // Use Lua script to atomically write server state
     let server_info_bytes = server_info.encode_to_vec();
     let keys = vec![server_key, workers_key];
-    let string_args = vec![ttl.as_secs().to_string()];
-    let binary_args = vec![server_info_bytes];
+    let mut args = vec![
+      RedisArg::Int(ttl.as_secs() as i64),
+      RedisArg::Bytes(server_info_bytes),
+    ];
 
     // ARGV[1] = TTL in seconds
     // ARGV[2] = server info (encoded protobuf as binary)
     // ARGV[3+] = worker info (暂时为空，需要支持 WorkerInfo)
     // ARGV[3+] = worker info (currently empty, needs to support WorkerInfo)
-
+    for worker in workers {
+      let worker_info = worker.encode_to_vec();
+      args.push(RedisArg::Str(worker.task_id));
+      args.push(RedisArg::Bytes(worker_info));
+    }
     let _: String = self
       .script_manager
-      .eval_script_with_binary_args(
-        &mut conn,
-        "write_server_state",
-        &keys,
-        &string_args,
-        &binary_args,
-      )
+      .eval_script(&mut conn, "write_server_state", &keys, &args)
       .await?;
 
     Ok(())
   }
 
   /// 清除服务器状态 - Go: ClearServerState
-  async fn clear_server_state(&self, host: &str, pid: i32, server_id: &str) -> Result<()> {
+  async fn clear_server_state(
+    &self,
+    host: &str,
+    pid: i32,
+    server_id: &str,
+    tenant: Option<&str>,
+  ) -> Result<()> {
     let mut conn = self.get_async_connection().await?;
 
     // 生成键
     // Generate keys
-    let server_key = keys::server_info_key(host, pid, server_id);
-    let workers_key = keys::workers_key(host, pid, server_id);
-
-    // 构造完整的 server_id (hostname:pid:uuid 格式)
-    // Construct full server_id (hostname:pid:uuid format)
-    let full_server_id = format!("{host}:{pid}:{server_id}");
-
+    let (server_key, workers_key) = keys::server_and_workers_keys(tenant, host, pid, server_id);
     // 1. 从 AllServers ZSET 中删除服务器 (服务器跟踪)
     // Remove the server from AllServers ZSET (server tracking)
-    let _: () = conn.zrem(keys::ALL_SERVERS, &full_server_id).await?;
+    let _: () = conn.zrem(keys::ALL_SERVERS, &server_key).await?;
 
     // 2. 从 AllWorkers ZSET 中删除工作者键 (如果存在)
     // Remove the worker key from AllWorkers ZSET (if exists)
@@ -888,18 +897,9 @@ impl Broker for RedisBroker {
     // 3. 使用 Lua 脚本原子性地清除服务器状态
     // Use Lua script to atomically clear server state
     let keys = vec![server_key, workers_key];
-    let string_args: Vec<String> = vec![];
-    let binary_args: Vec<Vec<u8>> = vec![];
-
     let _: String = self
       .script_manager
-      .eval_script_with_binary_args(
-        &mut conn,
-        "clear_server_state",
-        &keys,
-        &string_args,
-        &binary_args,
-      )
+      .eval_script(&mut conn, "clear_server_state", &keys, &[])
       .await?;
 
     Ok(())

asynq/src/backend/rdb/inspect.rs

@@ -990,7 +990,8 @@ impl RedisBroker {
 
     // 使用 clear_server_state 来正确地注销服务器
     // Use clear_server_state to correctly unregister the server
-    self.clear_server_state(hostname, pid, uuid).await
+    // Note: This is an administrative action, so no tenant isolation is applied
+    self.clear_server_state(hostname, pid, uuid, None).await
   }
 
   /// 心跳检测。

asynq/src/backend/rdb/redis_scripts.rs

@@ -1582,67 +1582,6 @@ impl ScriptManager {
   pub fn get_script_sha(&self, name: &str) -> Option<&String> {
     self.script_sha1.get(name)
   }
-
-  /// 执行脚本（支持二进制参数）
-  pub async fn eval_script_with_binary_args<T>(
-    &self,
-    conn: &mut RedisConnection,
-    script_name: &str,
-    keys: &[String],
-    string_args: &[String],
-    binary_args: &[Vec<u8>],
-  ) -> Result<T>
-  where
-    T: redis::FromRedisValue,
-  {
-    // 首先尝试使用 EVALSHA 如果脚本已加载
-    if let Some(sha) = self.get_script_sha(script_name) {
-      let mut cmd = redis::cmd("EVALSHA");
-      cmd.arg(sha).arg(keys.len()).arg(keys);
-
-      // 添加字符串参数
-      for arg in string_args {
-        cmd.arg(arg);
-      }
-
-      // 添加二进制参数
-      for arg in binary_args {
-        cmd.arg(arg);
-      }
-
-      match cmd.query_async::<T>(conn).await {
-        Ok(result) => return Ok(result),
-        Err(e) if e.to_string().contains("NOSCRIPT") => {
-          // 脚本被清理了，继续使用 EVAL
-        }
-        Err(e) => return Err(e.into()),
-      }
-    }
-
-    // 如果脚本未加载或 EVALSHA 失败，使用 EVAL
-    let script = match script_name {
-      "write_server_state" => scripts::WRITE_SERVER_STATE,
-      "clear_server_state" => scripts::CLEAR_SERVER_STATE,
-      _ => return Err(Error::other(format!("Script not loaded: {script_name}"))),
-    };
-
-    let mut cmd = redis::cmd("EVAL");
-    cmd.arg(script).arg(keys.len()).arg(keys);
-
-    // 添加字符串参数
-    for arg in string_args {
-      cmd.arg(arg);
-    }
-
-    // 添加二进制参数
-    for arg in binary_args {
-      cmd.arg(arg);
-    }
-
-    let result: T = cmd.query_async(conn).await?;
-    Ok(result)
-  }
-
   /// 执行脚本
   pub async fn eval_script<T>(
     &self,