tina-core 0.0.2

//! 基于Redis的限流实现

use crate::tina::core::service::throttle::IThrottleService;
use crate::tina::data::throttle::{Throttle, ThrottleParam};
use crate::tina::data::AppResult;
use crate::tina::redis::{IRedisClient, RedisConnection};
use crate::tina::server::session::Session;
use crate::{app_system_error, tina::data::throttle::ThrottleNewParam};
use indexmap::IndexMap;
use redis::AsyncCommands;
use std::cmp::max;

fn get_search_keys(params: &[ThrottleParam]) -> AppResult<Vec<String>> {
    let list = params.iter().map(|param| format!("throttle:{}:{}", param.key.key, param.key.duration_millis)).collect();
    Ok(list)
}

fn get_update_key(throttle: &Throttle) -> AppResult<String> {
    Ok(format!("throttle:{}:{}", throttle.key.key, throttle.key.duration_millis))
}

fn get_remove_keys(throttles: &[Throttle]) -> AppResult<Vec<String>> {
    let list = throttles.iter().map(|throttle| format!("throttle:{}:{}", throttle.key.key, throttle.key.duration_millis)).collect();
    Ok(list)
}

/// 基于Redis的限流实现
#[derive(Debug)]
pub(crate) struct RedisThrottleService;

#[allow(deprecated)]
#[async_trait]
impl IThrottleService for RedisThrottleService {
    #[instrument]
    async fn retrieve_throttle(&self, session: &Session, params: &[ThrottleParam]) -> AppResult<Vec<Throttle>> {
        let search_keys = get_search_keys(params)?;
        tracing::debug!("获取限流: params = {:?}", params);

        let application = session.get_application();
        let mut conn = application.get_redis_connection().await?;
        let cur_time = session.get_request_date_time().await?;

        let mut exists = Vec::with_capacity(search_keys.len());
        for search_key in search_keys.iter() {
            let exist: Option<Throttle> = conn.get(search_key.as_str()).await.map_err(crate::app_error_from!())?;
            if let Some(exist) = exist {
                exists.push(exist);
            }
        }

        let mut exist_map = IndexMap::new();
        for exist in exists.into_iter() {
            let key = exist.key.clone();
            exist_map.insert(key, exist);
        }

        let mut valid_map = IndexMap::new();
        let mut expires = Vec::new();
        for (key, value) in exist_map.into_iter() {
            match self.is_expired(cur_time, &value) {
                true => {
                    expires.push(value);
                }
                false => {
                    valid_map.insert(key, value);
                }
            }
        }

        if !expires.is_empty() {
            let remove_keys = get_remove_keys(&expires)?;
            for key in remove_keys.iter() {
                conn.del::<&str, ()>(key.as_str()).await.map_err(crate::app_error_from!())?;
            }
        }

        let mut list = Vec::with_capacity(params.len());
        for param in params.iter() {
            let key = &param.key;
            let exist = valid_map.remove(key);
            let new_param = ThrottleNewParam {
                current_date_time: cur_time,
                exist,
                key: key.key.to_owned(),
                threshold: param.threshold,
                duration_millis: key.duration_millis,
                fire_flag: false,
                param_map: param.param_map.clone(),
                message: param.message.clone(),
            };
            let new_throttle = self.build_new_throttle(new_param);
            list.push(new_throttle);
        }
        if list.is_empty() {
            return Err(app_system_error!("no throttle param retrieve from source: {:?}", params));
        }
        Ok(list)
    }

    #[instrument]
    async fn fire_throttle(&self, session: &Session, params: &[ThrottleParam]) -> AppResult<Vec<Throttle>> {
        let search_keys = get_search_keys(params)?;
        // tracing::debug!("检查限流: params = {}", params.to_json_string());

        let application = session.get_application();
        let mut conn = application.get_redis_connection().await?;
        let cur_time = session.get_request_date_time().await?;

        let mut exists = Vec::with_capacity(search_keys.len());
        for search_key in search_keys.iter() {
            let exist: Option<Throttle> = conn.get(search_key.as_str()).await.map_err(crate::app_error_from!())?;
            if let Some(exist) = exist {
                exists.push(exist);
            }
        }

        let mut exist_map = IndexMap::new();
        for exist in exists.into_iter() {
            let key = exist.key.clone();
            exist_map.insert(key, exist);
        }

        let mut valid_map = IndexMap::new();
        let mut expires = Vec::new();
        for (key, value) in exist_map.into_iter() {
            match self.is_expired(cur_time, &value) {
                true => {
                    expires.push(value);
                }
                false => {
                    valid_map.insert(key, value);
                }
            }
        }

        if !expires.is_empty() {
            let remove_keys = get_remove_keys(&expires)?;
            for key in remove_keys.iter() {
                conn.del::<&str, ()>(key.as_str()).await.map_err(crate::app_error_from!())?;
            }
        }

        let mut update_list = Vec::with_capacity(params.len());
        let mut update_duration_list = Vec::with_capacity(params.len());
        for param in params.iter() {
            let key = &param.key;
            let exist = valid_map.remove(key);
            let new_param = ThrottleNewParam {
                current_date_time: cur_time,
                exist,
                key: key.key.to_owned(),
                threshold: param.threshold,
                duration_millis: key.duration_millis,
                fire_flag: true,
                param_map: param.param_map.clone(),
                message: param.message.clone(),
            };
            let new_throttle = self.build_new_throttle(new_param);

            let update_key = get_update_key(&new_throttle)?;
            update_duration_list.push((update_key.to_owned(), key.duration_millis));
            update_list.push((update_key, new_throttle));
        }

        if !update_list.is_empty() && !update_duration_list.is_empty() {
            let mut pipe = crate::redis::pipe();
            let mut tmp_list = Vec::with_capacity(update_list.len());
            for item in update_list.iter() {
                tmp_list.push((&item.0, &item.1));
            }
            pipe.set_multiple(tmp_list.as_ref());
            for tuple in update_duration_list.iter() {
                let expire_millis = max(tuple.1, 0) as usize;
                pipe.pexpire(tuple.0.as_str(), expire_millis);
            }
            pipe.query_async::<RedisConnection, ()>(&mut conn).await.map_err(crate::app_error_from!())?;
            tracing::debug!("更新限流缓存: {:?}", tmp_list);
        }
        Ok(update_list.into_iter().map(|t| t.1).collect())
    }
}