use redis::{Script, aio::ConnectionManager};
use crate::{
BucketSize, HardLimitFactor, RateLimitComparator, SuppressionFactorCachePeriod, TrypemaError,
WindowSize,
common::{HistoryUpdateMode, RateType},
hybrid::{RedisProxyCommitter, common::StateRevision},
redis::{
RedisKey, RedisKeyGenerator,
scripts::{
SUPPRESSED_CLEANUP_LUA, SUPPRESSED_HYBRID_CLEAR_LUA,
SUPPRESSED_HYBRID_COMMIT_STATE_LUA, SUPPRESSED_HYBRID_DELETE_LUA,
SUPPRESSED_HYBRID_READ_STATE_LUA, SUPPRESSED_HYBRID_SET_RATE_LIMIT_LUA,
SUPPRESSED_SET_IF_LUA, lua_script, suppressed_lua_script,
},
},
};
#[derive(Debug)]
pub(crate) struct SuppressedHybridCommit {
pub key: RedisKey,
pub hard_window_limit: f64,
pub count: u64,
pub declined_count: u64,
pub state_revision: StateRevision,
}
#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct SuppressedHybridPendingState {
pub pending_count: u64,
pub pending_declined_count: u64,
pub state_revision: StateRevision,
}
#[derive(Debug)]
pub(crate) struct SuppressedHybridRedisProxyReadStateResult {
pub key: RedisKey,
pub current_total_count: u64,
pub current_total_declined: u64,
pub suppression_factor: f64,
pub suppression_factor_ttl_ms: Option<u64>,
pub hard_window_limit: Option<f64>,
pub state_revision: StateRevision,
}
#[derive(Clone, Debug)]
pub(crate) struct SuppressedHybridRedisProxy {
key_generator: RedisKeyGenerator,
read_state_script: Script,
commit_state_script: Script,
set_if_script: Script,
cleanup_script: Script,
set_rate_limit_script: Script,
delete_script: Script,
clear_script: Script,
connection_manager: ConnectionManager,
hard_limit_factor: HardLimitFactor,
suppression_factor_cache_period: SuppressionFactorCachePeriod,
bucket_size: BucketSize,
window_size: WindowSize,
read_chunk_size: usize,
}
pub(crate) struct SuppressedHybridRedisProxyOptions {
pub hard_limit_factor: HardLimitFactor,
pub suppression_factor_cache_period: SuppressionFactorCachePeriod,
pub bucket_size: BucketSize,
pub window_size: WindowSize,
pub prefix: RedisKey,
pub connection_manager: ConnectionManager,
}
impl SuppressedHybridRedisProxy {
pub(crate) fn new(
SuppressedHybridRedisProxyOptions {
prefix,
connection_manager,
hard_limit_factor,
suppression_factor_cache_period,
bucket_size,
window_size,
}: SuppressedHybridRedisProxyOptions,
) -> Self {
Self {
key_generator: RedisKeyGenerator::new(prefix, RateType::HybridSuppressed),
read_state_script: suppressed_lua_script(SUPPRESSED_HYBRID_READ_STATE_LUA),
commit_state_script: suppressed_lua_script(SUPPRESSED_HYBRID_COMMIT_STATE_LUA),
set_if_script: lua_script(SUPPRESSED_SET_IF_LUA),
cleanup_script: lua_script(SUPPRESSED_CLEANUP_LUA),
set_rate_limit_script: lua_script(SUPPRESSED_HYBRID_SET_RATE_LIMIT_LUA),
delete_script: suppressed_lua_script(SUPPRESSED_HYBRID_DELETE_LUA),
clear_script: lua_script(SUPPRESSED_HYBRID_CLEAR_LUA),
hard_limit_factor,
suppression_factor_cache_period,
connection_manager,
bucket_size,
window_size,
read_chunk_size: 100,
}
}
pub(crate) async fn read_state(
self: &SuppressedHybridRedisProxy,
key: &RedisKey,
) -> Result<SuppressedHybridRedisProxyReadStateResult, TrypemaError> {
let mut connection_manager = self.connection_manager.clone();
let res: (String, f64, u64, u64, f64, i64, u64, u64) = self
.read_state_script
.key(self.key_generator.get_hash_key(key))
.key(self.key_generator.get_active_keys(key))
.key(self.key_generator.get_window_limit_key(key))
.key(self.key_generator.get_total_count_key(key))
.key(self.key_generator.get_active_entities_key())
.key(self.key_generator.get_suppression_factor_key(key))
.key(self.key_generator.get_total_declined_key(key))
.key(self.key_generator.get_hash_declined_key(key))
.key(self.key_generator.get_state_revision_key())
.key(self.key_generator.get_key_state_revisions_key())
.arg(key.as_str())
.arg(self.window_size.as_seconds())
.arg(self.suppression_factor_cache_period.as_milliseconds())
.arg(self.hard_limit_factor.as_multiplier())
.invoke_async(&mut connection_manager)
.await?;
Ok(map_redis_read_result_to_state(res))
}
#[inline]
fn build_commit_pipeline(
&self,
commits: &[SuppressedHybridCommit],
should_load_script: bool,
) -> redis::Pipeline {
let mut pipe = redis::Pipeline::new();
if should_load_script {
pipe.load_script(&self.commit_state_script).ignore();
}
for commit in commits {
pipe.invoke_script(
self.commit_state_script
.key(self.key_generator.get_hash_key(&commit.key))
.key(self.key_generator.get_active_keys(&commit.key))
.key(self.key_generator.get_window_limit_key(&commit.key))
.key(self.key_generator.get_total_count_key(&commit.key))
.key(self.key_generator.get_active_entities_key())
.key(self.key_generator.get_suppression_factor_key(&commit.key))
.key(self.key_generator.get_total_declined_key(&commit.key))
.key(self.key_generator.get_hash_declined_key(&commit.key))
.key(self.key_generator.get_state_revision_key())
.key(self.key_generator.get_key_state_revisions_key())
.arg(commit.key.as_str())
.arg(self.window_size.as_seconds())
.arg(commit.hard_window_limit)
.arg(self.bucket_size.as_milliseconds())
.arg(self.suppression_factor_cache_period.as_milliseconds())
.arg(self.hard_limit_factor.as_multiplier())
.arg(commit.count)
.arg(commit.declined_count)
.arg(commit.state_revision.namespace)
.arg(commit.state_revision.key),
);
}
pipe
}
pub(crate) async fn batch_read_state(
self: &SuppressedHybridRedisProxy,
keys: &[RedisKey],
) -> Result<Vec<SuppressedHybridRedisProxyReadStateResult>, TrypemaError> {
if keys.is_empty() {
return Ok(Vec::new());
}
let mut connection_manager = self.connection_manager.clone();
let chunk_size = self.read_chunk_size.max(1);
let mut all_results: Vec<SuppressedHybridRedisProxyReadStateResult> =
Vec::with_capacity(keys.len());
for chunk in keys.chunks(chunk_size) {
let pipe = self.build_read_pipeline(chunk, false);
let results = match pipe
.query_async::<Vec<(String, f64, u64, u64, f64, i64, u64, u64)>>(
&mut connection_manager,
)
.await
{
Ok(results) => results,
Err(err) => {
if err.kind() != redis::ErrorKind::Server(redis::ServerErrorKind::NoScript) {
tracing::error!("redis.read.error, error executing pipeline: {:?}", err);
eprintln!("redis.read.error, error executing pipeline: {:?}", err);
return Err(TrypemaError::RedisError(err));
}
let pipe = self.build_read_pipeline(chunk, true);
match pipe
.query_async::<Vec<(String, f64, u64, u64, f64, i64, u64, u64)>>(
&mut connection_manager,
)
.await
{
Ok(results) => results,
Err(err) => {
tracing::error!(
"redis.read.error, error executing pipeline: {:?}",
err
);
eprintln!("redis.read.error, error executing pipeline: {:?}", err);
return Err(TrypemaError::RedisError(err));
}
}
}
};
all_results.extend(results.into_iter().map(map_redis_read_result_to_state));
}
Ok(all_results)
}
#[inline]
fn build_read_pipeline(&self, keys: &[RedisKey], should_load_script: bool) -> redis::Pipeline {
let mut pipe = redis::Pipeline::new();
if should_load_script {
pipe.load_script(&self.read_state_script).ignore();
}
for key in keys {
pipe.invoke_script(
self.read_state_script
.key(self.key_generator.get_hash_key(key))
.key(self.key_generator.get_active_keys(key))
.key(self.key_generator.get_window_limit_key(key))
.key(self.key_generator.get_total_count_key(key))
.key(self.key_generator.get_active_entities_key())
.key(self.key_generator.get_suppression_factor_key(key))
.key(self.key_generator.get_total_declined_key(key))
.key(self.key_generator.get_hash_declined_key(key))
.key(self.key_generator.get_state_revision_key())
.key(self.key_generator.get_key_state_revisions_key())
.arg(key.as_str())
.arg(self.window_size.as_seconds())
.arg(self.suppression_factor_cache_period.as_milliseconds())
.arg(self.hard_limit_factor.as_multiplier()),
);
}
pipe
}
pub(crate) async fn set_if(
&self,
key: &RedisKey,
hard_window_limit: f64,
comparator: RateLimitComparator,
count: u64,
mode: HistoryUpdateMode,
pending: SuppressedHybridPendingState,
) -> Result<(u64, u64, bool), TrypemaError> {
let mut connection_manager = self.connection_manager.clone();
let (comparator_op, comparator_operand) = comparator.redis_args();
let (new_total, old_total, changed): (u64, u64, u64) = self
.set_if_script
.key(self.key_generator.get_hash_key(key))
.key(self.key_generator.get_active_keys(key))
.key(self.key_generator.get_window_limit_key(key))
.key(self.key_generator.get_total_count_key(key))
.key(self.key_generator.get_active_entities_key())
.key(self.key_generator.get_suppression_factor_key(key))
.key(self.key_generator.get_total_declined_key(key))
.key(self.key_generator.get_hash_declined_key(key))
.arg(key.as_str())
.arg(self.window_size.as_seconds())
.arg(hard_window_limit)
.arg(comparator_op)
.arg(comparator_operand)
.arg(count)
.arg(mode.redis_arg())
.arg(pending.pending_count)
.arg(pending.pending_declined_count)
.invoke_async(&mut connection_manager)
.await?;
Ok((new_total, old_total, changed != 0))
}
pub(crate) async fn set_rate_limit(
&self,
key: &RedisKey,
hard_window_limit: f64,
) -> Result<Option<(f64, bool)>, TrypemaError> {
let mut connection_manager = self.connection_manager.clone();
let (status, previous, changed): (String, String, u8) = self
.set_rate_limit_script
.key(self.key_generator.get_window_limit_key(key))
.key(self.key_generator.get_active_entities_key())
.key(self.key_generator.get_suppression_factor_key(key))
.key(self.key_generator.get_key_state_revisions_key())
.arg(key.as_str())
.arg(self.window_size.as_seconds())
.arg(hard_window_limit)
.arg(self.hard_limit_factor.as_multiplier())
.invoke_async(&mut connection_manager)
.await?;
match status.as_str() {
"missing" => Ok(None),
"found" => previous
.parse::<f64>()
.map(|previous| Some((previous, changed == 1)))
.map_err(|_| {
TrypemaError::CustomError(
"invalid stored suppressed hybrid hard window limit".to_string(),
)
}),
"invalid" => Err(TrypemaError::CustomError(
"invalid stored suppressed hybrid hard window limit".to_string(),
)),
_ => Err(TrypemaError::UnexpectedRedisScriptResult {
operation: "suppressed_hybrid.set_rate_limit",
key: key.to_string(),
result: status,
}),
}
}
pub(crate) async fn delete(
&self,
key: &RedisKey,
pending_count: u64,
pending_declined_count: u64,
local_existed: bool,
) -> Result<Option<u64>, TrypemaError> {
let mut connection_manager = self.connection_manager.clone();
let (existed, accepted_count): (u8, u64) = self
.delete_script
.key(self.key_generator.get_hash_key(key))
.key(self.key_generator.get_hash_declined_key(key))
.key(self.key_generator.get_active_keys(key))
.key(self.key_generator.get_window_limit_key(key))
.key(self.key_generator.get_total_count_key(key))
.key(self.key_generator.get_total_declined_key(key))
.key(self.key_generator.get_suppression_factor_key(key))
.key(self.key_generator.get_active_entities_key())
.key(self.key_generator.get_key_state_revisions_key())
.arg(key.as_str())
.arg(self.window_size.as_seconds())
.arg(pending_count)
.arg(pending_declined_count)
.arg(u8::from(local_existed))
.invoke_async(&mut connection_manager)
.await?;
Ok((existed == 1).then_some(accepted_count))
}
pub(crate) async fn clear(&self) -> Result<(), TrypemaError> {
let mut connection_manager = self.connection_manager.clone();
let _: () = self
.clear_script
.key(self.key_generator.get_active_entities_key())
.key(self.key_generator.get_state_revision_key())
.key(self.key_generator.get_key_state_revisions_key())
.arg(self.key_generator.prefix.to_string())
.arg(self.key_generator.rate_type.to_string())
.arg(self.key_generator.hash_key_suffix.to_string())
.arg(self.key_generator.hash_declined_key_suffix.to_string())
.arg(self.key_generator.active_keys_key_suffix.to_string())
.arg(self.key_generator.window_limit_key_suffix.to_string())
.arg(self.key_generator.total_count_key_suffix.to_string())
.arg(self.key_generator.total_declined_key_suffix.to_string())
.arg(self.key_generator.suppression_factor_key_suffix.to_string())
.invoke_async(&mut connection_manager)
.await?;
Ok(())
}
pub(crate) async fn cleanup(&self, stale_after_ms: u64) -> Result<(), TrypemaError> {
let mut connection_manager = self.connection_manager.clone();
let _: () = self
.cleanup_script
.key(self.key_generator.prefix.to_string())
.key(self.key_generator.rate_type.to_string())
.key(self.key_generator.get_active_entities_key())
.arg(stale_after_ms)
.arg(self.key_generator.hash_key_suffix.to_string())
.arg(self.key_generator.window_limit_key_suffix.to_string())
.arg(self.key_generator.total_count_key_suffix.to_string())
.arg(self.key_generator.active_keys_key_suffix.to_string())
.arg(self.key_generator.suppression_factor_key_suffix.to_string())
.arg(self.key_generator.total_declined_key_suffix.to_string())
.arg(self.key_generator.hash_declined_key_suffix.to_string())
.invoke_async(&mut connection_manager)
.await?;
Ok(())
}
}
#[async_trait::async_trait]
impl RedisProxyCommitter<SuppressedHybridCommit> for SuppressedHybridRedisProxy {
async fn batch_commit_state(
self: &SuppressedHybridRedisProxy,
commits: &[SuppressedHybridCommit],
) -> Result<(), TrypemaError> {
let mut connection_manager = self.connection_manager.clone();
let pipe = self.build_commit_pipeline(commits, false);
let _: () = match pipe.query_async(&mut connection_manager).await {
Ok(results) => results,
Err(err) => {
if err.kind() != redis::ErrorKind::Server(redis::ServerErrorKind::NoScript) {
tracing::error!("redis.commit.error, error executing pipeline: {:?}", err);
return Err(TrypemaError::RedisError(err));
}
let pipe = self.build_commit_pipeline(commits, true);
match pipe.query_async::<()>(&mut connection_manager).await {
Ok(results) => results,
Err(err) => {
tracing::error!("redis.commit.error, error executing pipeline: {:?}", err);
return Err(TrypemaError::RedisError(err));
}
}
}
};
Ok(())
} }
fn map_redis_read_result_to_state(
(
entity,
suppression_factor,
current_total_count,
current_total_declined,
hard_window_limit,
suppression_factor_ttl_ms,
state_revision,
key_state_revision,
): (String, f64, u64, u64, f64, i64, u64, u64),
) -> SuppressedHybridRedisProxyReadStateResult {
SuppressedHybridRedisProxyReadStateResult {
key: RedisKey::from(entity),
current_total_count,
current_total_declined,
suppression_factor,
hard_window_limit: if hard_window_limit < 0.0 {
None
} else {
Some(hard_window_limit)
},
suppression_factor_ttl_ms: if suppression_factor_ttl_ms < 0 {
None
} else {
Some(suppression_factor_ttl_ms as u64)
},
state_revision: StateRevision {
namespace: state_revision,
key: key_state_revision,
},
}
}