feat: Add impl of #[dist_lock]

Author: itinycheng

Cargo.toml

@@ -7,12 +7,14 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [workspace]
+members = [".", "codegen"]
 
 [dependencies]
 chrono = "0.4.26"
 thiserror = "1.0.44"
 gethostname = "0.4.3"
 cfg-if = "1.0.0"
+dist_lock_codegen = { path = "./codegen", version = "*", optional = true }
 redis = { version = "0.23.0", optional = true }
 diesel = { version = "2.1.0", features = ["chrono"], optional = true }
 tokio = { version = "1.29.1", features = ["macros"], optional = true }
@@ -22,30 +24,32 @@ futures = { version = "0.3.28", optional = true }
 r2d2 = { version = "0.8.10", optional = true }
 
 [features]
-default = ["diesel_mysql_r2d2"]
+default = ["diesel_sqlite"]
 
 # redis
-redis_provider = ["redis/cluster"]
-redis_r2d2_provider = ["redis_provider", "redis/r2d2", "r2d2"]
+redis_provider = ["redis/cluster", "dist_lock_codegen/redis"]
+redis_r2d2_provider = ["redis_provider", "redis/r2d2", "r2d2", "dist_lock_codegen/redis"]
 redis_tokio_provider = [
     "redis/tokio-comp",
     "redis/cluster-async",
     "tokio",
     "futures",
     "async-trait",
+    "dist_lock_codegen/redis"
 ]
 redis_async_std_provider = [
     "redis/async-std-comp",
     "redis/cluster-async",
     "async-std",
     "futures",
     "async-trait",
+    "dist_lock_codegen/redis"
 ]
 
 # diesel
-diesel_sqlite = ["diesel/sqlite"]
-diesel_postgres = ["diesel/postgres"]
-diesel_mysql = ["diesel/mysql"]
-diesel_sqlite_r2d2 = ["diesel_sqlite", "diesel/r2d2", "r2d2"]
-diesel_postgres_r2d2 = ["diesel_postgres", "diesel/r2d2", "r2d2"]
-diesel_mysql_r2d2 = ["diesel_mysql", "diesel/r2d2", "r2d2"]
+diesel_sqlite = ["diesel/sqlite", "dist_lock_codegen/diesel"]
+diesel_postgres = ["diesel/postgres", "dist_lock_codegen/diesel"]
+diesel_mysql = ["diesel/mysql", "dist_lock_codegen/diesel"]
+diesel_sqlite_r2d2 = ["diesel_sqlite", "diesel/r2d2", "r2d2", "dist_lock_codegen/diesel"]
+diesel_postgres_r2d2 = ["diesel_postgres", "diesel/r2d2", "r2d2", "dist_lock_codegen/diesel"]
+diesel_mysql_r2d2 = ["diesel_mysql", "diesel/r2d2", "r2d2", "dist_lock_codegen/diesel"]

codegen/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "dist_lock_codegen"
+version = "0.0.1"
+authors = ["chenglong.wang <[email protected]>"]
+edition = "2021"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+proc-macro2 = "1.0.69"
+quote = "1.0.33"
+syn = { version = "2.0.38", features = ["full"] }
+humantime = "2.1.0"
+
+[features]
+redis = []
+diesel = []

codegen/src/impl_lock.rs

@@ -0,0 +1,73 @@
+use proc_macro2::Ident;
+use proc_macro2::Span;
+use proc_macro2::TokenStream;
+use quote::quote;
+use syn::ExprCall;
+use syn::Result;
+
+use crate::parse_args::DistLockArgs;
+
+pub(crate) fn generate(lock_args: DistLockArgs) -> Result<TokenStream> {
+	let name = lock_args.name.value();
+	let at_most_string = lock_args.at_most.value();
+	let at_most = at_most_string
+		.parse::<humantime::Duration>()
+		.map(|d| d.as_millis() as i64)
+		.map_err(|_| {
+			syn::Error::new(Span::call_site(), format!("can't prase at_most: {}", at_most_string))
+		})?;
+
+	let at_least_string = lock_args.at_least.map(|s| s.value()).unwrap_or("0s".to_string());
+	let at_least = at_least_string
+		.parse::<humantime::Duration>()
+		.map(|d| d.as_millis() as i64)
+		.map_err(|_| {
+			syn::Error::new(Span::call_site(), format!("can't prase at_least: {}", at_most_string))
+		})?;
+
+	Ok(gen_lock_code(name, at_most, at_least, lock_args.connection))
+}
+
+#[cfg(feature = "redis")]
+fn gen_lock_code(
+	name: String,
+	at_most_mills: i64,
+	at_least_mills: i64,
+	connection: ExprCall,
+) -> TokenStream {
+	let name_ident = Ident::new(&format!("_{}", &name), Span::call_site());
+	quote! {
+	   let mut #name_ident = {
+			use ::dist_lock::core::DistLock;
+			use ::dist_lock::core::LockConfig;
+			use ::dist_lock::provider::redis::RedisDriver;
+
+			let lock_name = #name.to_string();
+			let driver = RedisDriver::new(&lock_name, #connection);
+			let config = LockConfig::from_mills(lock_name, #at_least_mills, #at_most_mills);
+			DistLock::new(config, driver)
+		}
+	}
+}
+
+#[cfg(feature = "diesel")]
+fn gen_lock_code(
+	name: String,
+	at_most_mills: i64,
+	at_least_mills: i64,
+	connection: ExprCall,
+) -> TokenStream {
+	let name_ident = Ident::new(&format!("_{}", &name), Span::call_site());
+	quote! {
+	   let mut #name_ident = {
+			use ::dist_lock::core::DistLock;
+			use ::dist_lock::core::LockConfig;
+			use ::dist_lock::provider::diesel::DieselDriver;
+			
+			let lock_name = #name.to_string();
+			let driver = DieselDriver::new(&lock_name, Some("t"), #connection);
+			let config = LockConfig::from_mills(lock_name, #at_least_mills, #at_most_mills);
+			DistLock::new(config, driver)
+		}
+	}
+}

codegen/src/lib.rs

@@ -0,0 +1,38 @@
+use crate::parse_args::DistLockArgs;
+use impl_lock::generate;
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::parse_macro_input;
+use syn::ItemFn;
+use syn::Result;
+
+mod impl_lock;
+mod parse_args;
+
+#[proc_macro_attribute]
+pub fn dist_lock(head: TokenStream, body: TokenStream) -> TokenStream {
+	let lock_args = parse_macro_input!(head as DistLockArgs);
+	let function = parse_macro_input!(body as ItemFn);
+	parse(lock_args, function).unwrap_or_else(to_compile_error).into()
+}
+
+fn parse(lock_args: DistLockArgs, function: ItemFn) -> Result<proc_macro2::TokenStream> {
+	let dist_lock = generate(lock_args)?;
+	let fn_vis = function.vis;
+	let fn_body = function.block;
+	let fn_sig = function.sig;
+	let fn_name = fn_sig.ident;
+	let fn_generics = fn_sig.generics;
+	let fn_inputs = fn_sig.inputs;
+	let fn_output = fn_sig.output;
+	Ok(quote! {
+		#fn_vis fn #fn_name #fn_generics (#fn_inputs) #fn_output{
+			#dist_lock;
+			#fn_body
+		}
+	})
+}
+
+fn to_compile_error(e: syn::Error) -> proc_macro2::TokenStream {
+	e.to_compile_error()
+}

codegen/src/parse_args.rs

@@ -0,0 +1,64 @@
+use proc_macro2::Span;
+use syn::parenthesized;
+use syn::parse::Parse;
+use syn::ExprCall;
+use syn::LitStr;
+use syn::Token;
+
+pub(crate) struct DistLockArgs {
+	pub(crate) name: LitStr,
+	pub(crate) at_most: LitStr,
+	pub(crate) at_least: Option<LitStr>,
+	pub(crate) connection: ExprCall,
+}
+
+impl Parse for DistLockArgs {
+	fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+		let mut name = None;
+		let mut at_least = None;
+		let mut at_most = None;
+		let mut connection = None;
+		while !input.is_empty() {
+			let lookahead = input.lookahead1();
+			if lookahead.peek(kw::name) {
+				_ = input.parse::<kw::name>()?;
+				_ = input.parse::<Token![=]>()?;
+				name = Some(input.parse()?);
+			} else if lookahead.peek(kw::at_most) {
+				_ = input.parse::<kw::at_most>()?;
+				_ = input.parse::<Token![=]>()?;
+				at_most = Some(input.parse::<LitStr>()?);
+			} else if lookahead.peek(kw::at_least) {
+				_ = input.parse::<kw::at_least>()?;
+				_ = input.parse::<Token![=]>()?;
+				at_least = Some(input.parse::<LitStr>()?);
+			} else if lookahead.peek(kw::connection) {
+				_ = input.parse::<kw::connection>()?;
+				let content;
+				parenthesized!(content in input);
+				connection = Some(content.parse::<ExprCall>()?);
+			} else if lookahead.peek(Token![,]) {
+				_ = input.parse::<Token![,]>()?;
+			} else {
+				return Err(lookahead.error());
+			}
+		}
+
+		Ok(DistLockArgs {
+			name: name.ok_or(syn::Error::new(Span::call_site(), "lock name not found"))?,
+			at_most: at_most.ok_or(syn::Error::new(Span::call_site(), "at_most not found"))?,
+			at_least,
+			connection: connection
+				.ok_or(syn::Error::new(Span::call_site(), "connection not found"))?,
+		})
+	}
+}
+
+mod kw {
+	use syn::custom_keyword;
+
+	custom_keyword!(name);
+	custom_keyword!(at_least);
+	custom_keyword!(at_most);
+	custom_keyword!(connection);
+}

rust-toolchain.toml

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

src/core.rs

@@ -77,6 +77,10 @@ impl LockConfig {
 		LockConfig { name, min_lock, max_lock }
 	}
 
+	pub fn from_mills(name: String, min_lock: i64, max_lock: i64) -> LockConfig {
+		Self::new(name, Duration::milliseconds(min_lock), Duration::milliseconds(max_lock))
+	}
+
 	pub fn name(&self) -> &String {
 		&self.name
 	}

src/lib.rs

@@ -2,6 +2,8 @@ pub mod core;
 pub mod error;
 pub mod provider;
 
+pub use dist_lock_codegen::dist_lock;
+
 #[cfg(test)]
 mod tests {
 	use chrono::Duration;

src/provider/diesel.rs

@@ -3,7 +3,6 @@ use std::fmt::Display;
 use chrono::Utc;
 use diesel::sql_types::BigInt;
 use diesel::sql_types::VarChar;
-use diesel::QueryableByName;
 use diesel::RunQueryDsl;
 use gethostname::gethostname;
 
@@ -19,26 +18,6 @@ use super::help::sql_stmt::update_lock_sql;
 
 const LOCK_TABLE: &'static str = "dist_lock";
 
-diesel::table! {
-	dist_lock (name) {
-		#[max_length = 64]
-		name -> Varchar,
-		lock_until -> Bigint,
-		locked_at -> Bigint,
-		#[max_length = 255]
-		locked_by -> Varchar,
-	}
-}
-
-#[derive(QueryableByName)]
-#[diesel(table_name = crate::provider::diesel::dist_lock)]
-pub struct LockRecord {
-	pub name: String,
-	pub lock_until: i64,
-	pub locked_at: i64,
-	pub locked_by: String,
-}
-
 #[derive(Debug)]
 pub struct DieselDriver<T> {
 	name: String,

src/provider/help/mod.rs

@@ -1 +1,2 @@
+#[cfg(feature = "diesel")]
 pub(crate) mod sql_stmt;