first commit

Author: itinycheng

.gitignore

@@ -0,0 +1,4 @@
+/target
+/Cargo.lock
+/.idea/
+.DS_Store

Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "dist_lock"
+version = "0.0.1"
+authors = ["chenglong.wang <[email protected]>"]
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[workspace]
+
+[dependencies]
+chrono = "0.4.26"
+thiserror = "1.0.44"
+redis = { version = "0.23.0", optional = true }
+tokio = { version = "1.29.1", optional = true }
+async-trait = "0.1.72"
+
+[features]
+default = ["redis_tokio_provider"]
+redis_tokio_provider = ["redis/tokio-comp", "tokio/rt"]
+thread = []

rust-toolchain.toml

@@ -0,0 +1,3 @@
+[toolchain]
+channel = "nightly"
+components = ["rustfmt", "clippy", "rust-analyzer"]

rustfmt.toml

@@ -0,0 +1,9 @@
+hard_tabs = true
+max_width = 100
+use_small_heuristics = "Max"
+imports_granularity = "Item"
+reorder_imports = true
+# Format comments
+comment_width = 100
+wrap_comments = true
+edition = "2021"

src/core/local_cache.rs

@@ -0,0 +1,40 @@
+use chrono::DateTime;
+use chrono::Utc;
+
+#[cfg(feature = "thread")]
+use core::cell::Cell;
+
+use crate::provider::redis::RedisLock;
+
+#[cfg(feature = "thread")]
+thread_local! {
+	static LOCKED: Cell<DateTime<Utc>> = Cell::new(Default::default());
+}
+
+#[cfg(feature = "tokio")]
+tokio::task_local! {
+	static LOCKED: DateTime<Utc>;
+}
+
+pub fn set_local(time: DateTime<Utc>) {
+	#[cfg(feature = "thread")]
+	LOCKED.with(|cell| cell.set(time));
+
+	#[cfg(feature = "tokio")]
+	LOCKED.sync_scope(time, || {});
+}
+
+pub fn get_local() -> DateTime<Utc> {
+	#[cfg(feature = "thread")]
+	let value = LOCKED.with(|cell| cell.get());
+
+	#[cfg(feature = "tokio")]
+	let value = LOCKED.get();
+	value
+}
+
+impl<T> Drop for RedisLock<'_, T> {
+	fn drop(&mut self) {
+		set_local(Default::default());
+	}
+}

src/core/mod.rs

@@ -0,0 +1,22 @@
+use async_trait::async_trait;
+use chrono::Duration;
+
+use crate::error::LockResult;
+
+pub mod local_cache;
+
+#[derive(Debug)]
+pub struct LockConfig {
+	pub name: String,
+	pub min_lock: Duration,
+	pub max_lock: Duration,
+}
+
+#[async_trait]
+pub trait Lockable<T>: Sized {
+	async fn acquire(&mut self) -> LockResult<bool>;
+
+	async fn release(&self) -> LockResult<bool>;
+
+	async fn extend(&mut self) -> LockResult<bool>;
+}

src/error.rs

@@ -0,0 +1,10 @@
+pub type LockResult<T> = core::result::Result<T, crate::error::LockError>;
+
+#[derive(Debug, thiserror::Error)]
+pub enum LockError {
+	#[error("Redis error: {0}")]
+	RedisError(#[from] redis::RedisError),
+
+	#[error("lock failed")]
+	LockFailed,
+}

src/lib.rs

@@ -0,0 +1,15 @@
+pub mod core;
+pub mod error;
+pub mod provider;
+
+#[cfg(test)]
+mod tests {
+	use chrono::Duration;
+	use chrono::Utc;
+
+	#[test]
+	fn it_works() {
+		let a = Duration::max_value();
+		let _b = Utc::now() + a;
+	}
+}

src/provider/mod.rs

@@ -0,0 +1 @@
+pub mod redis;

src/provider/redis.rs

@@ -0,0 +1,83 @@
+use async_trait::async_trait;
+use chrono::DateTime;
+use chrono::Utc;
+use redis::Client;
+use redis::Value;
+
+use crate::core::LockConfig;
+use crate::core::Lockable;
+use crate::error::LockResult;
+
+const KEY_PREFIX: &str = "dist_lock";
+
+#[derive(Debug)]
+pub struct RedisLock<'a, T> {
+	pub key: String,
+	pub config: LockConfig,
+	pub driver: &'a T,
+	pub create_at: DateTime<Utc>,
+	pub locked_at: Option<DateTime<Utc>>,
+}
+
+impl<'a, T> RedisLock<'a, T> {
+	pub fn new(config: LockConfig, driver: &'a T) -> Self {
+		RedisLock {
+			key: format!("{}:{}", KEY_PREFIX, &config.name),
+			config,
+			driver,
+			create_at: Utc::now(),
+			locked_at: None,
+		}
+	}
+}
+
+#[async_trait]
+impl<'a> Lockable<Client> for RedisLock<'a, Client> {
+	async fn acquire(&mut self) -> LockResult<bool> {
+		let mut conn = self.driver.get_async_connection().await?;
+		let value: Value = redis::cmd("SET")
+			.arg(&self.key)
+			.arg(Utc::now().timestamp_millis())
+			.arg("NX")
+			.arg("PX")
+			.arg(self.config.max_lock.num_milliseconds() as usize)
+			.query_async(&mut conn)
+			.await?;
+
+		match value {
+			Value::Okay => {
+				self.locked_at = Some(Utc::now());
+				Ok(true)
+			}
+			_ => Ok(false),
+		}
+	}
+
+	async fn release(&self) -> LockResult<bool> {
+		let now = Utc::now();
+		let elapsed = now - self.locked_at.unwrap_or_default();
+		let remaining = self.config.min_lock - elapsed;
+		let mut conn = self.driver.get_async_connection().await?;
+		let value: Value = if remaining.num_milliseconds() > 0 {
+			redis::cmd("SET")
+				.arg(&self.key)
+				.arg(Utc::now().timestamp_millis())
+				.arg("XX")
+				.arg("PX")
+				.arg((self.config.min_lock - remaining).num_milliseconds())
+				.query_async(&mut conn)
+				.await?
+		} else {
+			redis::cmd("DEL").arg(&self.key).query_async(&mut conn).await?
+		};
+
+		match value {
+			Value::Okay => Ok(true),
+			_ => Ok(false),
+		}
+	}
+
+	async fn extend(&mut self) -> LockResult<bool> {
+		todo!()
+	}
+}