distributed_lock_postgres/
key.rs1use distributed_lock_core::error::{LockError, LockResult};
4use sha1::{Digest, Sha1};
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
11pub enum PostgresAdvisoryLockKey {
12 Single(i64),
14 Pair(i32, i32),
16}
17
18impl PostgresAdvisoryLockKey {
19 const MAX_ASCII_LENGTH: usize = 9;
21 const ASCII_CHAR_BITS: u32 = 7;
23 const MAX_ASCII_VALUE: u32 = (1 << Self::ASCII_CHAR_BITS) - 1;
25 const HASH_STRING_LENGTH: usize = 16;
27 const HASH_PART_LENGTH: usize = 8;
29 const HASH_STRING_SEPARATOR: char = ',';
31
32 pub fn from_name(name: &str, allow_hashing: bool) -> LockResult<Self> {
39 if name.is_empty() {
40 return Err(LockError::InvalidName(
41 "lock name cannot be empty".to_string(),
42 ));
43 }
44
45 if let Some(key) = Self::try_encode_ascii(name) {
47 return Ok(Self::Single(key));
48 }
49
50 if let Some(key) = Self::try_parse_hex_string(name) {
52 return Ok(key);
53 }
54
55 if let Some(key) = Self::try_parse_pair_string(name) {
57 return Ok(key);
58 }
59
60 if allow_hashing {
62 let hash = Self::hash_string(name);
63 return Ok(Self::Single(hash));
64 }
65
66 Err(LockError::InvalidName(format!(
67 "Name '{}' could not be encoded as a PostgresAdvisoryLockKey. Please specify allow_hashing or use one of the following formats: (1) a 0-{} character string using only ASCII characters, (2) a {} character hex string, or (3) a 2-part, {} character string of the form XXXXXXXX{}XXXXXXXX",
68 name,
69 Self::MAX_ASCII_LENGTH,
70 Self::HASH_STRING_LENGTH,
71 Self::HASH_PART_LENGTH * 2 + 1,
72 Self::HASH_STRING_SEPARATOR
73 )))
74 }
75
76 fn try_encode_ascii(name: &str) -> Option<i64> {
78 if name.len() > Self::MAX_ASCII_LENGTH {
79 return None;
80 }
81
82 let mut result = 0i64;
83 for ch in name.chars() {
84 let ch_val = ch as u32;
85 if ch_val > Self::MAX_ASCII_VALUE {
86 return None;
87 }
88 result = (result << Self::ASCII_CHAR_BITS) | (ch_val as i64);
89 }
90
91 result <<= 1;
93 for _ in name.len()..Self::MAX_ASCII_LENGTH {
94 result = (result << Self::ASCII_CHAR_BITS) | (Self::MAX_ASCII_VALUE as i64);
95 }
96
97 Some(result)
98 }
99
100 fn try_parse_hex_string(name: &str) -> Option<Self> {
102 if name.len() != Self::HASH_STRING_LENGTH {
103 return None;
104 }
105
106 i64::from_str_radix(name, 16).ok().map(Self::Single)
107 }
108
109 fn try_parse_pair_string(name: &str) -> Option<Self> {
111 let parts: Vec<&str> = name.split(Self::HASH_STRING_SEPARATOR).collect();
112 if parts.len() != 2 {
113 return None;
114 }
115
116 let key1 = i32::from_str_radix(parts[0], 16).ok()?;
117 let key2 = i32::from_str_radix(parts[1], 16).ok()?;
118
119 Some(Self::Pair(key1, key2))
120 }
121
122 #[allow(clippy::disallowed_methods)]
124 fn hash_string(name: &str) -> i64 {
125 let mut hasher = Sha1::new();
126 hasher.update(name.as_bytes());
127 let hash_bytes = hasher.finalize();
128
129 let mut result = 0i64;
133 for i in (0..8).rev() {
134 result = (result << 8) | (hash_bytes[i] as i64);
135 }
136 result
137 }
138
139 pub fn has_single_key(&self) -> bool {
141 matches!(self, Self::Single(_))
142 }
143
144 pub fn key(&self) -> i64 {
146 match self {
147 Self::Single(k) => *k,
148 Self::Pair(_, _) => panic!("key() called on Pair variant"),
149 }
150 }
151
152 pub fn keys(&self) -> (i32, i32) {
154 match self {
155 Self::Single(k) => {
156 let upper = (*k >> 32) as i32;
158 let lower = (*k & 0xFFFFFFFF) as i32;
159 (upper, lower)
160 }
161 Self::Pair(k1, k2) => (*k1, *k2),
162 }
163 }
164
165 #[allow(clippy::wrong_self_convention)]
167 pub fn to_sql_args(&self) -> String {
168 match self {
169 Self::Single(k) => format!("{k}"),
170 Self::Pair(k1, k2) => format!("{k1}, {k2}"),
171 }
172 }
173}
174
175#[cfg(test)]
176mod tests {
177 use super::*;
178
179 #[test]
180 fn test_ascii_encoding() {
181 let key = PostgresAdvisoryLockKey::from_name("test", false).unwrap();
182 assert!(key.has_single_key());
183 }
184
185 #[test]
186 fn test_hex_encoding() {
187 let hex_str = "0000000000000001";
188 let key = PostgresAdvisoryLockKey::from_name(hex_str, false).unwrap();
189 assert!(key.has_single_key());
190 assert_eq!(key.key(), 1);
191 }
192
193 #[test]
194 fn test_pair_encoding() {
195 let pair_str = "00000001,00000002";
196 let key = PostgresAdvisoryLockKey::from_name(pair_str, false).unwrap();
197 match key {
198 PostgresAdvisoryLockKey::Pair(k1, k2) => {
199 assert_eq!(k1, 1);
200 assert_eq!(k2, 2);
201 }
202 _ => panic!("Expected Pair variant"),
203 }
204 }
205
206 #[test]
207 fn test_hash_encoding() {
208 let long_name = "this is a very long lock name that needs hashing";
209 let key = PostgresAdvisoryLockKey::from_name(long_name, true).unwrap();
210 assert!(key.has_single_key());
211 }
212}