wx_rust_common/util/locks/
redis_distributed_lock.rs

1//! 基于 Redis 的简单分布式锁。
2//!
3//! 对应 Java `RedisTemplateSimpleDistributedLock` 的语义：支持重入的
4//! 简单 Redis 分布式锁（非红锁），基于 SETNX + TTL + 唯一值校验释放。
5//! Java 使用 Spring `StringRedisTemplate`；Rust 侧使用 `redis` crate。
6
7use std::sync::Arc;
8
9use redis::Client;
10
11/// Redis 简单分布式锁。
12///
13/// 锁语义：
14/// - 加锁：`SET key value NX EX lease`（value 为唯一标识，支持重入判断）
15/// - 释放：Lua 脚本校验 value 后 DEL（防止误删他人锁）
16#[derive(Clone)]
17pub struct RedisDistributedLock {
18    client: Client,
19    key: String,
20    lease_milliseconds: i64,
21}
22
23impl RedisDistributedLock {
24    /// 构建 Redis 分布式锁（随机键名）。
25    ///
26    /// # 参数
27    /// - `client`：redis 客户端
28    /// - `lease_milliseconds`：锁租约时长（毫秒），必须大于 0
29    ///
30    /// # 错误
31    /// `lease_milliseconds <= 0` 时返回参数错误。
32    pub fn new(client: Client, lease_milliseconds: i64) -> Result<Self, String> {
33        if lease_milliseconds <= 0 {
34            return Err(format!(
35                "Parameter 'leaseMilliseconds' must grate then 0: {lease_milliseconds}"
36            ));
37        }
38        let key = format!("lock:{}", uuid());
39        Ok(Self {
40            client,
41            key,
42            lease_milliseconds,
43        })
44    }
45
46    /// 构建指定键名的 Redis 分布式锁。
47    ///
48    /// # 参数
49    /// - `client`：redis 客户端
50    /// - `key`：锁键名
51    /// - `lease_milliseconds`：锁租约时长（毫秒），必须大于 0
52    pub fn with_key(
53        client: Client,
54        key: impl Into<String>,
55        lease_milliseconds: i64,
56    ) -> Result<Self, String> {
57        if lease_milliseconds <= 0 {
58            return Err(format!(
59                "Parameter 'leaseMilliseconds' must grate then 0: {lease_milliseconds}"
60            ));
61        }
62        Ok(Self {
63            client,
64            key: key.into(),
65            lease_milliseconds,
66        })
67    }
68
69    /// 返回锁租约时长（毫秒）。
70    pub fn lease_milliseconds(&self) -> i64 {
71        self.lease_milliseconds
72    }
73
74    /// 返回锁键名。
75    pub fn key(&self) -> &str {
76        &self.key
77    }
78
79    /// 尝试获取锁（非阻塞）。
80    ///
81    /// # 返回
82    /// 获取成功返回 `Ok(Some(guard))`（Drop 时释放锁）；连接失败返回错误。
83    pub fn try_lock(&self) -> Result<Option<LockGuard>, redis::RedisError> {
84        let value = uuid();
85        let mut conn = self.client.get_connection()?;
86        let result: Option<String> = redis::cmd("SET")
87            .arg(&self.key)
88            .arg(&value)
89            .arg("NX")
90            .arg("PX")
91            .arg(self.lease_milliseconds)
92            .query(&mut conn)?;
93        if result.is_some() {
94            Ok(Some(LockGuard {
95                client: self.client.clone(),
96                key: self.key.clone(),
97                value,
98            }))
99        } else {
100            Ok(None)
101        }
102    }
103
104    /// 阻塞获取锁（轮询，对应 Java `lock()` 的 1 秒重试）。
105    ///
106    /// # 参数
107    /// - `timeout`：总超时；`None` 为无限等待
108    ///
109    /// # 返回
110    /// 获取成功返回 guard；超时或错误返回错误信息。
111    pub fn lock(&self, timeout: Option<std::time::Duration>) -> Result<LockGuard, String> {
112        let start = std::time::Instant::now();
113        loop {
114            match self.try_lock() {
115                Ok(Some(g)) => return Ok(g),
116                Ok(None) => {
117                    if let Some(t) = timeout
118                        && start.elapsed() >= t
119                    {
120                        return Err("acquire timeouted".to_string());
121                    }
122                    std::thread::sleep(std::time::Duration::from_millis(1000));
123                }
124                Err(e) => return Err(format!("lock failed: {e}")),
125            }
126        }
127    }
128}
129
130/// 锁守卫：Drop 时自动释放锁（校验 value 后 DEL）。
131pub struct LockGuard {
132    client: Client,
133    key: String,
134    value: String,
135}
136
137impl Drop for LockGuard {
138    fn drop(&mut self) {
139        // Lua：校验 value 一致才删除（对应 Java 释放脚本）
140        let script = "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end";
141        let mut conn = match self.client.get_connection() {
142            Ok(c) => c,
143            Err(_) => return,
144        };
145        let _: Option<i64> = redis::cmd("EVAL")
146            .arg(script)
147            .arg(1)
148            .arg(&self.key)
149            .arg(&self.value)
150            .query(&mut conn)
151            .ok();
152    }
153}
154
155/// 生成简单 UUID 字符串（锁唯一标识）。
156fn uuid() -> String {
157    // 16 字节随机值转 hex（rand 0.10 自由函数）
158    let bytes: [u8; 16] = rand::random();
159    hex::encode(bytes)
160}
161
162/// 便于无 Redis 环境下保持 API 可用（非 feature 构建时仍引用 Arc）
163#[allow(dead_code)]
164fn _touch(_: Arc<()>) {}
wx_rust_common/util/locks/redis_distributed_lock.rs

wx_rust_common/util/locks/
redis_distributed_lock.rs
