tina-core 0.0.2

//! 基于Redis的分布式锁
#![allow(dead_code)]

use crate::tina::core::service::lock::IDistributedLockService;
use crate::tina::data::AppResult;
use crate::tina::redis::{PooledRedisClient, RedisConnection};
use crate::tina::util::not_empty::INotEmpty;
use futures_util::future;
use futures_util::future::Either;
use redis::AsyncCommands;
use std::cmp::min;
use std::future::Future;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, SystemTime};
use tracing::Instrument;

const DEFAULT_LOCK_CACHE_SECONDS: u64 = 30;
const DEFAULT_FLUSH_EXPIRE_INTERVAL_SECONDS: u64 = 5;

fn get_redis_key(key: &str, value: &str) -> String {
    format!("lock_object:{}:{}", key, value)
}

async fn drop_lock(key: &str, value: &str, client: PooledRedisClient) -> AppResult<()> {
    match key.not_empty() {
        true => {
            let redis_key = get_redis_key(key, value);
            let mut conn = client.get().await.map_err(crate::app_error_from!())?;
            conn.del::<'_, &str, ()>(redis_key.as_str()).await.map_err(crate::app_error_from!())?;
        }
        false => {}
    }
    tracing::debug!("分布式锁 {} - {}, 释放成功!", key, value);
    Ok(())
}

async fn sleep_current(duration: Duration) {
    tokio::time::sleep(duration).await
}

#[derive(Debug)]
pub struct RedisDistributedLock(pub(crate) PooledRedisClient, pub(crate) Option<(String, String)>);

#[crate::async_trait]
impl IDistributedLockService for RedisDistributedLock {
    async fn lock<Lock, F, R>(
        &mut self,
        key: &str,
        value: Option<Lock>,
        lock_timeout: Duration,
        retry_interval: Duration,
        fut: F,
    ) -> AppResult<R>
    where
        Lock: ToString + Send + 'static,
        F: Future<Output = AppResult<R>> + Send + 'static,
        R: Send + 'static,
    {
        let lock_value = match value.as_ref() {
            None => "".to_string(),
            Some(v) => v.to_string(),
        };
        let lock_timeout_secs = lock_timeout.as_secs();
        let retry_interval_secs = retry_interval.as_secs();
        let span = debug_span!("redis distributed lock", lock_value, lock_timeout_secs, retry_interval_secs);
        async move {
            match value {
                None => {
                    tracing::debug!("分布式锁 {} - {}, 不执行锁操作...", key, "None");
                    fut.await
                }
                Some(_) => {
                    let value = format!("{:x}", md5::compute(lock_value));
                    let value = value.as_str();
                    let redis_key = get_redis_key(key, value);
                    let start_time = SystemTime::now();
                    let mut duration;
                    let mut lock_result;
                    let mut retry_flag = false;
                    loop {
                        duration = SystemTime::now().duration_since(start_time).ok().unwrap_or_else(|| Duration::from_nanos(u64::MAX));
                        if duration >= lock_timeout {
                            tracing::error!("分布式锁 {} - {}, 获取锁超时...超时临界值: {}秒!", key, value, lock_timeout.as_secs());
                            return Err(crate::app_system_error!(
                                "分布式锁 {} - {}, 获取锁超时...超时临界值: {}秒!",
                                key,
                                value,
                                lock_timeout.as_secs()
                            ));
                        }
                        if retry_flag {
                            tracing::info!("分布式锁 {} - {}, 尝试获取锁...超时临界值: {}秒!", key, value, lock_timeout.as_secs());
                        }
                        retry_flag = true;
                        let script = r##"if (redis.call('exists', KEYS[1]) == 0) then
                            redis.call('hincrby', KEYS[1], ARGV[2], 1);
                            redis.call('expire', KEYS[1], ARGV[1]);
                            return nil;
                            end;
                            if (redis.call('hexists', KEYS[1], ARGV[2]) == 1) then
                            redis.call('hincrby', KEYS[1], ARGV[2], 1);
                            redis.call('expire', KEYS[1], ARGV[1]);
                            return nil;
                            end;
                            return redis.call('pttl', KEYS[1]);"##;
                        let script = crate::redis::Script::new(script);
                        let mut conn = self.0.get().await.map_err(crate::app_error_from!())?;
                        lock_result = script
                            .key(redis_key.as_str())
                            .arg(DEFAULT_LOCK_CACHE_SECONDS)
                            .arg(uuid::Uuid::new_v4().to_string())
                            .invoke_async::<RedisConnection, Option<i64>>(&mut conn)
                            .await
                            .map_err(crate::app_error_from!());
                        match lock_result {
                            Ok(ttl) => {
                                if ttl.unwrap_or(0) > 0 {
                                    sleep_current(retry_interval).await;
                                    continue;
                                }
                                self.1 = Some((key.to_owned(), value.to_owned()));
                                tracing::debug!("分布式锁 {} - {}, 获取锁成功...超时临界值: {}秒!", key, value, lock_timeout.as_secs());
                                break;
                            }
                            Err(_) => {
                                sleep_current(retry_interval).await;
                                continue;
                            }
                        }
                    }

                    let complete_flag = AtomicBool::new(false);
                    let job_fut = async {
                        let r = fut.await;
                        complete_flag.store(true, Ordering::Relaxed);
                        r
                    };
                    let flush_expire_fut = async {
                        loop {
                            if complete_flag.load(Ordering::Relaxed) {
                                break;
                            }
                            match self.0.get().await {
                                Ok(mut conn) => {
                                    if let Err(err) = conn.expire::<&str, ()>(redis_key.as_str(), DEFAULT_LOCK_CACHE_SECONDS as usize).await
                                    {
                                        tracing::error!("分布式锁 {} - {}, 刷新超时异常: {:?}", key, value, err);
                                    }
                                }
                                Err(err) => {
                                    tracing::error!("分布式锁 {} - {}, 获取刷新超时连接异常: {:?}", key, value, err);
                                }
                            }
                            sleep_current(min(Duration::from_secs(DEFAULT_FLUSH_EXPIRE_INTERVAL_SECONDS), retry_interval)).await;
                        }
                    };
                    futures_util::pin_mut!(job_fut);
                    futures_util::pin_mut!(flush_expire_fut);
                    let select1 = future::select(job_fut, flush_expire_fut).await;
                    let r = match select1 {
                        Either::Left((left, _)) => left,
                        Either::Right((_, right)) => right.await,
                    };
                    match drop_lock(key, value, self.0.to_owned()).await {
                        Ok(_) => r,
                        Err(err) => Err(err),
                    }
                }
            }
        }
        .instrument(span)
        .await
    }
}