mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
refactor: extract runtime state backends
This commit is contained in:
fawney19
@@ -4,25 +4,32 @@ use axum::http::Response;
|
||||
use futures_util::StreamExt;
|
||||
|
||||
use crate::concurrency::ConcurrencyPermit;
|
||||
use crate::distributed::DistributedConcurrencyPermit;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct AdmissionPermit {
|
||||
_local: Option<ConcurrencyPermit>,
|
||||
_distributed: Option<DistributedConcurrencyPermit>,
|
||||
_distributed: Option<Box<dyn Send + Sync>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for AdmissionPermit {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("AdmissionPermit")
|
||||
.field("has_local", &self._local.is_some())
|
||||
.field("has_distributed", &self._distributed.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl AdmissionPermit {
|
||||
pub fn from_parts(
|
||||
pub fn from_parts<D: Send + Sync + 'static>(
|
||||
local: Option<ConcurrencyPermit>,
|
||||
distributed: Option<DistributedConcurrencyPermit>,
|
||||
distributed: Option<D>,
|
||||
) -> Option<Self> {
|
||||
if local.is_none() && distributed.is_none() {
|
||||
None
|
||||
} else {
|
||||
Some(Self {
|
||||
_local: local,
|
||||
_distributed: distributed,
|
||||
_distributed: distributed.map(|permit| Box::new(permit) as Box<dyn Send + Sync>),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -70,7 +77,7 @@ fn hold_axum_response_permit(response: Response<Body>, permit: AdmissionPermit)
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{hold_admission_permit_until, maybe_hold_axum_response_permit, AdmissionPermit};
|
||||
use crate::{ConcurrencyGate, DistributedConcurrencyGate};
|
||||
use crate::ConcurrencyGate;
|
||||
use axum::body::{to_bytes, Body};
|
||||
use axum::http::Response;
|
||||
|
||||
@@ -96,12 +103,9 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn holds_combined_local_and_distributed_permit_until_future_finishes() {
|
||||
let local_gate = ConcurrencyGate::new("local", 1);
|
||||
let distributed_gate = DistributedConcurrencyGate::new_in_memory("distributed", 1);
|
||||
let local = local_gate.try_acquire().expect("local permit");
|
||||
let distributed = distributed_gate
|
||||
.try_acquire()
|
||||
.await
|
||||
.expect("distributed permit");
|
||||
let distributed_gate = ConcurrencyGate::new("distributed", 1);
|
||||
let distributed = distributed_gate.try_acquire().expect("distributed permit");
|
||||
|
||||
let task = tokio::spawn(hold_admission_permit_until(
|
||||
AdmissionPermit::from_parts(Some(local), Some(distributed)),
|
||||
@@ -116,19 +120,12 @@ mod tests {
|
||||
"local permit should still be held"
|
||||
);
|
||||
assert!(
|
||||
distributed_gate.try_acquire().await.is_err(),
|
||||
distributed_gate.try_acquire().is_err(),
|
||||
"distributed permit should still be held"
|
||||
);
|
||||
|
||||
task.await.expect("task should complete");
|
||||
assert_eq!(local_gate.snapshot().in_flight, 0);
|
||||
assert_eq!(
|
||||
distributed_gate
|
||||
.snapshot()
|
||||
.await
|
||||
.expect("snapshot should build")
|
||||
.in_flight,
|
||||
0
|
||||
);
|
||||
assert_eq!(distributed_gate.snapshot().in_flight, 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,10 +1,4 @@
|
||||
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::concurrency::{ConcurrencyGate, ConcurrencyPermit};
|
||||
use crate::metrics::{MetricKind, MetricLabel, MetricSample};
|
||||
@@ -68,80 +62,11 @@ impl DistributedConcurrencySnapshot {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RedisDistributedConcurrencyConfig {
|
||||
pub url: String,
|
||||
pub key_prefix: Option<String>,
|
||||
pub lease_ttl_ms: u64,
|
||||
pub renew_interval_ms: u64,
|
||||
pub command_timeout_ms: Option<u64>,
|
||||
}
|
||||
|
||||
impl Default for RedisDistributedConcurrencyConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
url: String::new(),
|
||||
key_prefix: None,
|
||||
lease_ttl_ms: 30_000,
|
||||
renew_interval_ms: 10_000,
|
||||
command_timeout_ms: Some(1_000),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisDistributedConcurrencyConfig {
|
||||
fn validate(&self) -> Result<(), DistributedConcurrencyError> {
|
||||
let raw = self.url.trim();
|
||||
if raw.is_empty() {
|
||||
return Err(DistributedConcurrencyError::InvalidConfiguration(
|
||||
"distributed concurrency redis url cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
url::Url::parse(raw).map_err(|err| {
|
||||
DistributedConcurrencyError::InvalidConfiguration(format!(
|
||||
"invalid distributed concurrency redis url: {err}"
|
||||
))
|
||||
})?;
|
||||
if self.lease_ttl_ms == 0 {
|
||||
return Err(DistributedConcurrencyError::InvalidConfiguration(
|
||||
"distributed concurrency lease_ttl_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.renew_interval_ms == 0 {
|
||||
return Err(DistributedConcurrencyError::InvalidConfiguration(
|
||||
"distributed concurrency renew_interval_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.renew_interval_ms >= self.lease_ttl_ms {
|
||||
return Err(DistributedConcurrencyError::InvalidConfiguration(
|
||||
"distributed concurrency renew_interval_ms must be smaller than lease_ttl_ms"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
if matches!(self.command_timeout_ms, Some(0)) {
|
||||
return Err(DistributedConcurrencyError::InvalidConfiguration(
|
||||
"distributed concurrency command_timeout_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn semaphore_key(&self, gate: &'static str) -> String {
|
||||
prefixed_key(self.key_prefix.as_deref(), &format!("admission:{gate}"))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum DistributedConcurrencyBackend {
|
||||
InMemory(Arc<ConcurrencyGate>),
|
||||
Redis(Arc<RedisDistributedState>),
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct DistributedConcurrencyState {
|
||||
gate: &'static str,
|
||||
limit: usize,
|
||||
backend: DistributedConcurrencyBackend,
|
||||
gate_impl: Arc<ConcurrencyGate>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -159,49 +84,11 @@ impl DistributedConcurrencyGate {
|
||||
state: Arc::new(DistributedConcurrencyState {
|
||||
gate,
|
||||
limit,
|
||||
backend: DistributedConcurrencyBackend::InMemory(Arc::new(ConcurrencyGate::new(
|
||||
gate, limit,
|
||||
))),
|
||||
gate_impl: Arc::new(ConcurrencyGate::new(gate, limit)),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new_redis(
|
||||
gate: &'static str,
|
||||
limit: usize,
|
||||
config: RedisDistributedConcurrencyConfig,
|
||||
) -> Result<Self, DistributedConcurrencyError> {
|
||||
if limit == 0 {
|
||||
return Err(DistributedConcurrencyError::InvalidConfiguration(
|
||||
"distributed concurrency gate limit must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
config.validate()?;
|
||||
let client = redis::Client::open(config.url.clone()).map_err(|err| {
|
||||
DistributedConcurrencyError::InvalidConfiguration(format!(
|
||||
"failed to build distributed concurrency redis client: {err}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(Self {
|
||||
state: Arc::new(DistributedConcurrencyState {
|
||||
gate,
|
||||
limit,
|
||||
backend: DistributedConcurrencyBackend::Redis(Arc::new(RedisDistributedState {
|
||||
gate,
|
||||
limit,
|
||||
client,
|
||||
key: config.semaphore_key(gate),
|
||||
lease_ttl_ms: config.lease_ttl_ms,
|
||||
renew_interval_ms: config.renew_interval_ms,
|
||||
command_timeout_ms: config.command_timeout_ms,
|
||||
high_watermark: AtomicUsize::new(0),
|
||||
rejected: AtomicU64::new(0),
|
||||
})),
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn gate(&self) -> &'static str {
|
||||
self.state.gate
|
||||
}
|
||||
@@ -213,364 +100,46 @@ impl DistributedConcurrencyGate {
|
||||
pub async fn try_acquire(
|
||||
&self,
|
||||
) -> Result<DistributedConcurrencyPermit, DistributedConcurrencyError> {
|
||||
match &self.state.backend {
|
||||
DistributedConcurrencyBackend::InMemory(gate) => gate
|
||||
.try_acquire()
|
||||
.map(DistributedConcurrencyPermit::from_in_memory)
|
||||
.map_err(|err| match err {
|
||||
crate::ConcurrencyError::Saturated { gate, limit } => {
|
||||
DistributedConcurrencyError::Saturated { gate, limit }
|
||||
self.state
|
||||
.gate_impl
|
||||
.try_acquire()
|
||||
.map(|permit| DistributedConcurrencyPermit { _permit: permit })
|
||||
.map_err(|err| match err {
|
||||
crate::ConcurrencyError::Saturated { gate, limit } => {
|
||||
DistributedConcurrencyError::Saturated { gate, limit }
|
||||
}
|
||||
crate::ConcurrencyError::Closed { gate } => {
|
||||
DistributedConcurrencyError::Unavailable {
|
||||
gate,
|
||||
limit: self.state.limit,
|
||||
message: "in-memory distributed concurrency gate is closed".to_string(),
|
||||
}
|
||||
crate::ConcurrencyError::Closed { gate } => {
|
||||
DistributedConcurrencyError::Unavailable {
|
||||
gate,
|
||||
limit: self.state.limit,
|
||||
message: "in-memory distributed concurrency gate is closed".to_string(),
|
||||
}
|
||||
}
|
||||
}),
|
||||
DistributedConcurrencyBackend::Redis(state) => state.try_acquire().await,
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn snapshot(
|
||||
&self,
|
||||
) -> Result<DistributedConcurrencySnapshot, DistributedConcurrencyError> {
|
||||
match &self.state.backend {
|
||||
DistributedConcurrencyBackend::InMemory(gate) => {
|
||||
let snapshot = gate.snapshot();
|
||||
Ok(DistributedConcurrencySnapshot {
|
||||
limit: snapshot.limit,
|
||||
in_flight: snapshot.in_flight,
|
||||
available_permits: snapshot.available_permits,
|
||||
high_watermark: snapshot.high_watermark,
|
||||
rejected: snapshot.rejected,
|
||||
})
|
||||
}
|
||||
DistributedConcurrencyBackend::Redis(state) => state.snapshot().await,
|
||||
}
|
||||
let snapshot = self.state.gate_impl.snapshot();
|
||||
Ok(DistributedConcurrencySnapshot {
|
||||
limit: snapshot.limit,
|
||||
in_flight: snapshot.in_flight,
|
||||
available_permits: snapshot.available_permits,
|
||||
high_watermark: snapshot.high_watermark,
|
||||
rejected: snapshot.rejected,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct DistributedConcurrencyPermit {
|
||||
inner: DistributedConcurrencyPermitInner,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum DistributedConcurrencyPermitInner {
|
||||
InMemory(ConcurrencyPermit),
|
||||
Redis {
|
||||
state: Arc<RedisDistributedState>,
|
||||
token: String,
|
||||
renew_task: JoinHandle<()>,
|
||||
},
|
||||
}
|
||||
|
||||
impl DistributedConcurrencyPermit {
|
||||
fn from_in_memory(permit: ConcurrencyPermit) -> Self {
|
||||
Self {
|
||||
inner: DistributedConcurrencyPermitInner::InMemory(permit),
|
||||
}
|
||||
}
|
||||
|
||||
fn from_redis(
|
||||
state: Arc<RedisDistributedState>,
|
||||
token: String,
|
||||
renew_task: JoinHandle<()>,
|
||||
) -> Self {
|
||||
Self {
|
||||
inner: DistributedConcurrencyPermitInner::Redis {
|
||||
state,
|
||||
token,
|
||||
renew_task,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for DistributedConcurrencyPermit {
|
||||
fn drop(&mut self) {
|
||||
match &mut self.inner {
|
||||
DistributedConcurrencyPermitInner::InMemory(_permit) => {}
|
||||
DistributedConcurrencyPermitInner::Redis {
|
||||
state,
|
||||
token,
|
||||
renew_task,
|
||||
} => {
|
||||
renew_task.abort();
|
||||
let state = Arc::clone(state);
|
||||
let token = token.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(err) = state.release(&token).await {
|
||||
warn!(
|
||||
gate = state.gate,
|
||||
error = %err,
|
||||
"failed to release distributed concurrency permit"
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct RedisDistributedState {
|
||||
gate: &'static str,
|
||||
limit: usize,
|
||||
client: redis::Client,
|
||||
key: String,
|
||||
lease_ttl_ms: u64,
|
||||
renew_interval_ms: u64,
|
||||
command_timeout_ms: Option<u64>,
|
||||
high_watermark: AtomicUsize,
|
||||
rejected: AtomicU64,
|
||||
}
|
||||
|
||||
impl RedisDistributedState {
|
||||
async fn try_acquire(
|
||||
self: &Arc<Self>,
|
||||
) -> Result<DistributedConcurrencyPermit, DistributedConcurrencyError> {
|
||||
let token = format!("{}:{}", self.gate, Uuid::new_v4());
|
||||
let now_ms = unix_time_ms();
|
||||
let expires_at_ms = now_ms.saturating_add(self.lease_ttl_ms);
|
||||
let key = self.key.clone();
|
||||
let result: (i64, i64) = self
|
||||
.run_with_timeout("acquire", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_err(|err| self.unavailable(format!("connect failed: {err}")))?;
|
||||
redis::Script::new(
|
||||
"redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]); \
|
||||
local count = redis.call('ZCARD', KEYS[1]); \
|
||||
if count >= tonumber(ARGV[3]) then \
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[5]); \
|
||||
return {0, count}; \
|
||||
end; \
|
||||
redis.call('ZADD', KEYS[1], ARGV[2], ARGV[4]); \
|
||||
count = redis.call('ZCARD', KEYS[1]); \
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[5]); \
|
||||
return {1, count};",
|
||||
)
|
||||
.key(&key)
|
||||
.arg(now_ms as i64)
|
||||
.arg(expires_at_ms as i64)
|
||||
.arg(self.limit as i64)
|
||||
.arg(&token)
|
||||
.arg(self.lease_ttl_ms as i64)
|
||||
.invoke_async::<(i64, i64)>(&mut connection)
|
||||
.await
|
||||
.map_err(|err| self.unavailable(format!("acquire failed: {err}")))
|
||||
})
|
||||
.await?;
|
||||
|
||||
let acquired = result.0 > 0;
|
||||
let in_flight = result.1.max(0) as usize;
|
||||
self.observe_in_flight(in_flight);
|
||||
|
||||
if !acquired {
|
||||
self.rejected.fetch_add(1, Ordering::Relaxed);
|
||||
return Err(DistributedConcurrencyError::Saturated {
|
||||
gate: self.gate,
|
||||
limit: self.limit,
|
||||
});
|
||||
}
|
||||
|
||||
let renew_state = Arc::clone(self);
|
||||
let renew_token = token.clone();
|
||||
let renew_task = tokio::spawn(async move {
|
||||
let interval = Duration::from_millis(renew_state.renew_interval_ms);
|
||||
loop {
|
||||
tokio::time::sleep(interval).await;
|
||||
if let Err(err) = renew_state.renew(&renew_token).await {
|
||||
warn!(
|
||||
gate = renew_state.gate,
|
||||
error = %err,
|
||||
"failed to renew distributed concurrency permit"
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(DistributedConcurrencyPermit::from_redis(
|
||||
Arc::clone(self),
|
||||
token,
|
||||
renew_task,
|
||||
))
|
||||
}
|
||||
|
||||
async fn snapshot(
|
||||
&self,
|
||||
) -> Result<DistributedConcurrencySnapshot, DistributedConcurrencyError> {
|
||||
let in_flight = self.live_count().await?;
|
||||
Ok(DistributedConcurrencySnapshot {
|
||||
limit: self.limit,
|
||||
in_flight,
|
||||
available_permits: self.limit.saturating_sub(in_flight),
|
||||
high_watermark: self.high_watermark.load(Ordering::Relaxed),
|
||||
rejected: self.rejected.load(Ordering::Relaxed),
|
||||
})
|
||||
}
|
||||
|
||||
async fn renew(&self, token: &str) -> Result<(), DistributedConcurrencyError> {
|
||||
let now_ms = unix_time_ms();
|
||||
let expires_at_ms = now_ms.saturating_add(self.lease_ttl_ms);
|
||||
let key = self.key.clone();
|
||||
let renewed = self
|
||||
.run_with_timeout("renew", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_err(|err| self.unavailable(format!("connect failed: {err}")))?;
|
||||
redis::Script::new(
|
||||
"redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]); \
|
||||
local score = redis.call('ZSCORE', KEYS[1], ARGV[2]); \
|
||||
if not score then \
|
||||
return 0; \
|
||||
end; \
|
||||
redis.call('ZADD', KEYS[1], 'XX', ARGV[3], ARGV[2]); \
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[4]); \
|
||||
return 1;",
|
||||
)
|
||||
.key(&key)
|
||||
.arg(now_ms as i64)
|
||||
.arg(token)
|
||||
.arg(expires_at_ms as i64)
|
||||
.arg(self.lease_ttl_ms as i64)
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
.map_err(|err| self.unavailable(format!("renew failed: {err}")))
|
||||
})
|
||||
.await?;
|
||||
|
||||
if renewed == 0 {
|
||||
return Err(self.unavailable("lease token expired".to_string()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn release(&self, token: &str) -> Result<(), DistributedConcurrencyError> {
|
||||
let key = self.key.clone();
|
||||
self.run_with_timeout("release", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_err(|err| self.unavailable(format!("connect failed: {err}")))?;
|
||||
redis::Script::new(
|
||||
"local removed = redis.call('ZREM', KEYS[1], ARGV[1]); \
|
||||
if removed > 0 and redis.call('ZCARD', KEYS[1]) == 0 then \
|
||||
redis.call('DEL', KEYS[1]); \
|
||||
end; \
|
||||
return removed;",
|
||||
)
|
||||
.key(&key)
|
||||
.arg(token)
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
.map_err(|err| self.unavailable(format!("release failed: {err}")))?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn live_count(&self) -> Result<usize, DistributedConcurrencyError> {
|
||||
let now_ms = unix_time_ms();
|
||||
let key = self.key.clone();
|
||||
let count = self
|
||||
.run_with_timeout("snapshot", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_err(|err| self.unavailable(format!("connect failed: {err}")))?;
|
||||
redis::Script::new(
|
||||
"redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]); \
|
||||
return redis.call('ZCARD', KEYS[1]);",
|
||||
)
|
||||
.key(&key)
|
||||
.arg(now_ms as i64)
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
.map_err(|err| self.unavailable(format!("snapshot failed: {err}")))
|
||||
})
|
||||
.await?
|
||||
.max(0) as usize;
|
||||
self.observe_in_flight(count);
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
async fn run_with_timeout<T, F>(
|
||||
&self,
|
||||
operation: &'static str,
|
||||
future: F,
|
||||
) -> Result<T, DistributedConcurrencyError>
|
||||
where
|
||||
F: std::future::Future<Output = Result<T, DistributedConcurrencyError>>,
|
||||
{
|
||||
if let Some(timeout_ms) = self.command_timeout_ms {
|
||||
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
self.unavailable(format!(
|
||||
"{operation} exceeded {timeout_ms}ms command timeout"
|
||||
))
|
||||
})?
|
||||
} else {
|
||||
future.await
|
||||
}
|
||||
}
|
||||
|
||||
fn unavailable(&self, message: String) -> DistributedConcurrencyError {
|
||||
DistributedConcurrencyError::Unavailable {
|
||||
gate: self.gate,
|
||||
limit: self.limit,
|
||||
message,
|
||||
}
|
||||
}
|
||||
|
||||
fn observe_in_flight(&self, in_flight: usize) {
|
||||
let mut observed = self.high_watermark.load(Ordering::Acquire);
|
||||
while in_flight > observed {
|
||||
match self.high_watermark.compare_exchange_weak(
|
||||
observed,
|
||||
in_flight,
|
||||
Ordering::AcqRel,
|
||||
Ordering::Acquire,
|
||||
) {
|
||||
Ok(_) => break,
|
||||
Err(next) => observed = next,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn prefixed_key(prefix: Option<&str>, raw_key: &str) -> String {
|
||||
let prefix = prefix.unwrap_or_default().trim().trim_matches(':');
|
||||
if prefix.is_empty() {
|
||||
raw_key.trim_matches(':').to_string()
|
||||
} else {
|
||||
format!("{prefix}:{}", raw_key.trim_matches(':'))
|
||||
}
|
||||
}
|
||||
|
||||
fn unix_time_ms() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_millis() as u64
|
||||
_permit: ConcurrencyPermit,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
DistributedConcurrencyError, DistributedConcurrencyGate, RedisDistributedConcurrencyConfig,
|
||||
};
|
||||
use super::{DistributedConcurrencyError, DistributedConcurrencyGate};
|
||||
|
||||
#[tokio::test]
|
||||
async fn shared_in_memory_gate_rejects_second_acquire() {
|
||||
@@ -599,46 +168,4 @@ mod tests {
|
||||
let snapshot = gate.snapshot().await.expect("snapshot should build");
|
||||
assert_eq!(snapshot.in_flight, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_redis_config() {
|
||||
let error = DistributedConcurrencyGate::new_redis(
|
||||
"shared",
|
||||
1,
|
||||
RedisDistributedConcurrencyConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
lease_ttl_ms: 10_000,
|
||||
renew_interval_ms: 10_000,
|
||||
command_timeout_ms: Some(1_000),
|
||||
},
|
||||
)
|
||||
.expect_err("equal renew interval should fail");
|
||||
assert_eq!(
|
||||
error,
|
||||
DistributedConcurrencyError::InvalidConfiguration(
|
||||
"distributed concurrency renew_interval_ms must be smaller than lease_ttl_ms"
|
||||
.to_string()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_redis_gate_without_touching_network() {
|
||||
let gate = DistributedConcurrencyGate::new_redis(
|
||||
"gateway_requests_distributed",
|
||||
2,
|
||||
RedisDistributedConcurrencyConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
lease_ttl_ms: 15_000,
|
||||
renew_interval_ms: 5_000,
|
||||
command_timeout_ms: Some(1_000),
|
||||
},
|
||||
)
|
||||
.expect("redis gate should build");
|
||||
|
||||
assert_eq!(gate.gate(), "gateway_requests_distributed");
|
||||
assert_eq!(gate.limit(), 2);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -20,7 +20,7 @@ pub use concurrency::{ConcurrencyError, ConcurrencyGate, ConcurrencyPermit, Conc
|
||||
pub use config::ServiceRuntimeConfig;
|
||||
pub use distributed::{
|
||||
DistributedConcurrencyError, DistributedConcurrencyGate, DistributedConcurrencyPermit,
|
||||
DistributedConcurrencySnapshot, RedisDistributedConcurrencyConfig,
|
||||
DistributedConcurrencySnapshot,
|
||||
};
|
||||
pub use error::RuntimeBootstrapError;
|
||||
pub use metrics::{prometheus_response, service_up_sample, MetricKind, MetricLabel, MetricSample};
|
||||
|
||||
Reference in New Issue
Block a user