use crate::tina::data::i18n_string::I18nString;
use crate::tina::data::throttle::{Throttle, ThrottleFirePolicy, ThrottleParam, ThrottleUniqueKey};
use crate::tina::data::AppResult;
use crate::tina::data::{
date_time::{DurationFormat, LocalDateTime},
throttle::ThrottleNewParam,
};
use crate::tina::i18n::message::system_message::SystemMessage;
use crate::tina::server::session::Session;
use crate::tina::util::not_empty::INotEmpty;
use crate::{app_system_error, app_system_error_with_msg, i18n_string};
use chrono::Duration;
use std::cmp::{max, Ordering};
use std::fmt::Debug;
#[async_trait]
pub trait IThrottleService: Debug + Send + Sync + 'static {
async fn retrieve_throttle(&self, session: &Session, params: &[ThrottleParam]) -> AppResult<Vec<Throttle>>;
async fn fire_throttle(&self, session: &Session, params: &[ThrottleParam]) -> AppResult<Vec<Throttle>>;
#[instrument]
fn get_min_throttle(&self, throttles: Vec<Throttle>) -> Option<Throttle> {
let mut min_throttle: Option<Throttle> = None;
for throttle in throttles.into_iter() {
if min_throttle.is_none() {
min_throttle = Some(throttle);
} else if let Some(throttle2) = min_throttle.as_ref() {
let remains = throttle.threshold.unwrap_or(0) as i32 - throttle.fires.unwrap_or(0) as i32;
let min_remains = throttle2.threshold.unwrap_or(0) as i32 - throttle2.fires.unwrap_or(0) as i32;
match remains.cmp(&min_remains) {
Ordering::Less => {
min_throttle = Some(throttle);
}
Ordering::Equal => {
let expire_at = throttle.expire_at.as_ref();
let min_expire_at = throttle2.expire_at.as_ref();
if let Some(expire_at) = expire_at {
if let Some(min_expire_at) = min_expire_at {
if expire_at.0 > min_expire_at.0 {
min_throttle = Some(throttle);
}
}
}
}
Ordering::Greater => {}
}
}
}
min_throttle
}
#[instrument]
async fn handle_throttle(&self, session: &Session, params: &[ThrottleParam]) -> AppResult<Option<Throttle>> {
if params.is_empty() {
return Ok(None);
}
let throttles = self.retrieve_throttle(session, params).await?;
if throttles.is_empty() {
return Err(app_system_error!("获取限流剩余次数失败! params = {:?}", params));
}
let min_throttle = self.get_min_throttle(throttles);
if let Some(throttle) = min_throttle.as_ref() {
let remains = throttle.threshold.unwrap_or(0) as i32 - throttle.fires.unwrap_or(0) as i32;
if remains <= 0 {
let mut message = match throttle.message.as_ref() {
None => i18n_string!(SystemMessage::ERROR_REQUEST_TOO_FREQUENTLY),
Some(v) => v.clone(),
};
let req_time = session.get_request_date_time().await?;
fill_message_args(session, &mut message, throttle, req_time);
return Err(app_system_error_with_msg!(message, "无可用剩余限流次数! throttle = {:?}", min_throttle));
}
}
Ok(min_throttle)
}
#[instrument]
async fn after_handle_throttle_success(&self, session: &Session, params: &[ThrottleParam]) -> AppResult<()> {
if params.is_empty() {
return Ok(());
}
let fire_params = params
.iter()
.filter(|¶m| {
let fire_policy = param.fire_policy;
match fire_policy {
ThrottleFirePolicy::Always => true,
ThrottleFirePolicy::WithoutError => true,
ThrottleFirePolicy::WithError => false,
}
})
.cloned()
.collect::<Vec<ThrottleParam>>();
if fire_params.is_empty() {
return Ok(());
}
self.fire_throttle(session, &fire_params).await?;
Ok(())
}
#[instrument]
async fn after_handle_throttle_error(&self, session: &Session, params: &[ThrottleParam]) -> AppResult<()> {
if params.is_empty() {
return Ok(());
}
let fire_params: Vec<ThrottleParam> = params
.iter()
.filter(|¶m| {
let fire_policy = param.fire_policy;
match fire_policy {
ThrottleFirePolicy::Always => true,
ThrottleFirePolicy::WithoutError => false,
ThrottleFirePolicy::WithError => true,
}
})
.cloned()
.collect();
if fire_params.is_empty() {
return Ok(());
}
self.fire_throttle(session, &fire_params).await?;
Ok(())
}
fn build_new_throttle(&self, param: ThrottleNewParam) -> Throttle {
let ThrottleNewParam {
current_date_time,
exist,
key,
threshold,
duration_millis,
fire_flag,
param_map,
message,
} = param;
let mut new_throttle = match exist {
None => Throttle {
key: ThrottleUniqueKey {
key,
duration_millis,
},
fires: Some(0),
threshold: Some(threshold),
create_time: Some(current_date_time),
update_time: Some(current_date_time),
expire_at: None,
param_map,
message,
},
Some(exist) => Throttle {
key: exist.key,
fires: exist.fires,
threshold: exist.threshold,
create_time: exist.create_time,
update_time: exist.update_time,
expire_at: exist.expire_at,
param_map,
message,
},
};
new_throttle.key.duration_millis = duration_millis;
if fire_flag {
new_throttle.fires = Some(new_throttle.fires.unwrap_or(0) + 1);
}
new_throttle.threshold = Some(threshold);
if new_throttle.create_time.is_none() {
new_throttle.create_time = Some(current_date_time);
}
new_throttle.expire_at =
Some(LocalDateTime(new_throttle.create_time.unwrap_or(current_date_time).0 + Duration::milliseconds(duration_millis as i64)));
new_throttle.update_time = Some(current_date_time);
new_throttle
}
#[instrument]
fn is_expired(&self, cur_time: LocalDateTime, throttle: &Throttle) -> bool {
match throttle.expire_at.as_ref() {
None => true,
Some(expire_at) => *cur_time > **expire_at,
}
}
}
fn fill_message_args(session: &Session, message: &mut I18nString, throttle: &Throttle, req_time: LocalDateTime) {
for (name, value) in throttle.param_map.iter() {
if name.not_empty() && value.not_empty() {
message.add_arg(name, value);
}
}
if let Some(expire_at) = throttle.expire_at {
let duration = *expire_at - *req_time;
let remain_seconds = max(duration.num_seconds(), 0);
message.add_arg("seconds", remain_seconds.to_string());
message.add_arg("duration", duration.to_i18n_format().get_string(session.get_locale()));
}
}