69f12aa
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use export::{AppRedis, RedisConnection};
use uuid::Uuid;
use crate::smart::error::{SmartError, SmartResult};
#[async_trait]
pub trait PushOperations: Send + Sync {
async fn acquire(&self, key: &str, member: &[u8], ttl_secs: u64) -> SmartResult<bool>;
async fn renew(&self, key: &str, token: &[u8], ttl_secs: u64) -> SmartResult<bool>;
async fn release(&self, key: &str, token: &[u8]) -> SmartResult<()>;
async fn enqueue(&self, key: &str, member: &[u8]) -> SmartResult<()>;
async fn dequeue(&self, key: &str, member: &[u8]) -> SmartResult<()>;
}
macro_rules! dispatch {
($conn:expr, |$c:ident| $body:expr) => {
match &mut $conn {
RedisConnection::Standalone($c) => $body,
RedisConnection::Cluster($c) => $body,
}
};
}
const RENEW_SCRIPT: &str = "if redis.call('GET', KEYS[1]) == ARGV[1] then return redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) else return 0 end";
const RELEASE_SCRIPT: &str = "if redis.call('GET', KEYS[1]) == ARGV[1] then return redis.call('DEL', KEYS[1]) else return 0 end";
pub struct RedisPushBackend {
pub kv: Arc<AppRedis>,
}
impl RedisPushBackend {
async fn eval_bool(&self, script: &str, key: &str, argv: Vec<Vec<u8>>) -> SmartResult<bool> {
let mut conn = self.kv.connection_clone();
let mut cmd = redis::cmd("EVAL");
cmd.arg(script).arg(1).arg(key);
for a in argv {
cmd.arg(a);
}
let n: i64 = dispatch!(conn, |c| cmd.query_async(c).await)?;
Ok(n != 0)
}
}
#[async_trait]
impl PushOperations for RedisPushBackend {
async fn acquire(&self, key: &str, member: &[u8], ttl_secs: u64) -> SmartResult<bool> {
let mut conn = self.kv.connection_clone();
let resp: Option<String> = dispatch!(conn, |c| {
redis::cmd("SET")
.arg(key)
.arg(member)
.arg("NX")
.arg("EX")
.arg(ttl_secs)
.query_async(c)
.await
})?;
Ok(resp.is_some())
}
async fn renew(&self, key: &str, token: &[u8], ttl_secs: u64) -> SmartResult<bool> {
self.eval_bool(
RENEW_SCRIPT,
key,
vec![token.to_vec(), ttl_secs.to_string().into_bytes()],
)
.await
}
async fn release(&self, key: &str, token: &[u8]) -> SmartResult<()> {
self.eval_bool(RELEASE_SCRIPT, key, vec![token.to_vec()])
.await?;
Ok(())
}
async fn enqueue(&self, key: &str, member: &[u8]) -> SmartResult<()> {
let mut conn = self.kv.connection_clone();
let _: i64 = dispatch!(conn, |c| {
redis::cmd("RPUSH")
.arg(key)
.arg(member)
.query_async(c)
.await
})?;
Ok(())
}
async fn dequeue(&self, key: &str, member: &[u8]) -> SmartResult<()> {
let mut conn = self.kv.connection_clone();
let _: i64 = dispatch!(conn, |c| {
redis::cmd("LREM")
.arg(key)
.arg(1)
.arg(member)
.query_async(c)
.await
})?;
Ok(())
}
}
pub struct NoopPushBackend;
const NOPUSH_BACKEND_ERROR: &str =
"push requires redis: SmartService kv is not configured (NoopPushBackend)";
#[async_trait]
impl PushOperations for NoopPushBackend {
async fn acquire(&self, _key: &str, _member: &[u8], _ttl_secs: u64) -> SmartResult<bool> {
Err(SmartError::Internal(NOPUSH_BACKEND_ERROR.to_string()))
}
async fn renew(&self, _key: &str, _token: &[u8], _ttl_secs: u64) -> SmartResult<bool> {
Err(SmartError::Internal(NOPUSH_BACKEND_ERROR.to_string()))
}
async fn release(&self, _key: &str, _token: &[u8]) -> SmartResult<()> {
Ok(())
}
async fn enqueue(&self, _key: &str, _member: &[u8]) -> SmartResult<()> {
Err(SmartError::Internal(NOPUSH_BACKEND_ERROR.to_string()))
}
async fn dequeue(&self, _key: &str, _member: &[u8]) -> SmartResult<()> {
Err(SmartError::Internal(NOPUSH_BACKEND_ERROR.to_string()))
}
}
#[derive(Clone)]
pub struct PushCoordinator {
backend: Arc<dyn PushOperations>,
lock_ttl_secs: u64,
renew_interval_secs: u64,
poll_interval: Duration,
}
const DEFAULT_LOCK_TTL_SECS: u64 = 120;
const DEFAULT_RENEW_INTERVAL_SECS: u64 = 40;
const DEFAULT_POLL_INTERVAL: Duration = Duration::from_millis(200);
impl PushCoordinator {
pub fn new(backend: Arc<dyn PushOperations>) -> Self {
Self {
backend,
lock_ttl_secs: DEFAULT_LOCK_TTL_SECS,
renew_interval_secs: DEFAULT_RENEW_INTERVAL_SECS,
poll_interval: DEFAULT_POLL_INTERVAL,
}
}
pub async fn run<T>(
&self,
key: &str,
f: impl std::future::Future<Output = SmartResult<T>>,
) -> SmartResult<T> {
let token = Uuid::now_v7().to_string();
let result: SmartResult<T> = async {
if self
.backend
.acquire(key, token.as_bytes(), self.lock_ttl_secs)
.await?
{
tracing::debug!(key, ttl_secs = self.lock_ttl_secs, "push lock acquired");
return self.run_locked(key, &token, f).await;
}
tracing::debug!(key, "push lock busy; queued, waiting for lock");
self.backend.enqueue(key, token.as_bytes()).await?;
loop {
tokio::time::sleep(self.poll_interval).await;
if self
.backend
.acquire(key, token.as_bytes(), self.lock_ttl_secs)
.await?
{
self.backend.dequeue(key, token.as_bytes()).await?;
tracing::debug!(key, "push lock acquired after waiting");
return self.run_locked(key, &token, f).await;
}
}
}
.await;
result
}
async fn run_locked<T>(
&self,
key: &str,
token: &str,
f: impl std::future::Future<Output = SmartResult<T>>,
) -> SmartResult<T> {
let heartbeat = Heartbeat::spawn(
self.backend.clone(),
key.to_string(),
token.to_string(),
self.lock_ttl_secs,
self.renew_interval_secs,
);
let result = f.await;
heartbeat.stop().await;
self.backend.release(key, token.as_bytes()).await?;
tracing::debug!(key, ok = result.is_ok(), "push lock released");
result
}
}
struct Heartbeat {
handle: Option<tokio::task::JoinHandle<()>>,
stop: Arc<tokio::sync::Notify>,
}
impl Drop for Heartbeat {
fn drop(&mut self) {
self.stop.notify_one();
if let Some(handle) = &self.handle
&& !handle.is_finished()
{
handle.abort();
}
}
}
impl Heartbeat {
fn spawn(
backend: Arc<dyn PushOperations>,
key: String,
token: String,
ttl: u64,
interval: u64,
) -> Self {
let stop = Arc::new(tokio::sync::Notify::new());
let stop_clone = stop.clone();
let handle = tokio::spawn(async move {
loop {
tokio::select! {
_ = stop_clone.notified() => break,
_ = tokio::time::sleep(Duration::from_secs(interval)) => {
match backend.renew(&key, token.as_bytes(), ttl).await {
Ok(false) => {
tracing::debug!(key = %key, "push lock no longer held; stopping renewer");
break;
}
Err(error) => {
tracing::warn!(%error, "lock renew failed, stopping renewer");
break;
}
Ok(true) => {}
}
}
}
}
});
Self {
handle: Some(handle),
stop,
}
}
async fn stop(mut self) {
self.stop.notify_one();
if let Some(handle) = self.handle.take() {
let _ = handle.await;
}
}
}