1use crate::{env::RedisMode, response::ApiResponse};
2use axum::http::StatusCode;
3use compact_str::ToCompactString;
4use rustis::{
5 client::Client,
6 commands::{
7 GenericCommands, InfoSection, ServerCommands, SetCondition, SetExpiration, StringCommands,
8 },
9 resp::BulkString,
10};
11use serde::{Serialize, de::DeserializeOwned};
12use std::{
13 future::Future,
14 sync::{
15 Arc,
16 atomic::{AtomicU64, Ordering},
17 },
18 time::{Duration, Instant},
19};
20
21#[derive(Clone, Serialize)]
22pub struct BulkStringRef<'a>(
23 #[serde(
24 deserialize_with = "::rustis::resp::deserialize_byte_buf",
25 serialize_with = "::rustis::resp::serialize_byte_buf"
26 )]
27 pub &'a [u8],
28);
29
30#[derive(Clone, Debug)]
31struct DataEntry {
32 data: Arc<Vec<u8>>,
33 intended_ttl: Duration,
34}
35
36#[derive(Clone, Debug)]
37struct LockEntry {
38 semaphore: Arc<tokio::sync::Semaphore>,
39}
40
41struct DataExpiry;
42
43impl moka::Expiry<compact_str::CompactString, DataEntry> for DataExpiry {
44 fn expire_after_create(
45 &self,
46 _key: &compact_str::CompactString,
47 value: &DataEntry,
48 _created_at: Instant,
49 ) -> Option<Duration> {
50 Some(value.intended_ttl)
51 }
52}
53
54pub struct Cache {
55 client: Option<Arc<Client>>,
56 use_internal_cache: bool,
57 local: moka::future::Cache<compact_str::CompactString, DataEntry>,
58 local_task: tokio::task::JoinHandle<()>,
59 local_locks: moka::future::Cache<compact_str::CompactString, LockEntry>,
60 local_locks_task: tokio::task::JoinHandle<()>,
61 local_ratelimits: moka::future::Cache<compact_str::CompactString, (u64, u64)>,
62
63 cache_calls: AtomicU64,
64 cache_latency_ns_total: AtomicU64,
65 cache_latency_ns_max: AtomicU64,
66 cache_misses: AtomicU64,
67}
68
69impl Cache {
70 pub async fn new(env: &crate::env::Env) -> Arc<Self> {
71 let start = std::time::Instant::now();
72
73 let client = match &env.redis_mode {
74 RedisMode::Redis { redis_url } => {
75 if let Some(redis_url) = redis_url {
76 Some(Arc::new(Client::connect(redis_url.clone()).await.unwrap()))
77 } else {
78 None
79 }
80 }
81 RedisMode::Sentinel {
82 cluster_name,
83 redis_sentinels,
84 } => Some(Arc::new(
85 Client::connect(
86 format!(
87 "redis-sentinel://{}/{cluster_name}/0",
88 redis_sentinels.join(",")
89 )
90 .as_str(),
91 )
92 .await
93 .unwrap(),
94 )),
95 };
96
97 let local = moka::future::Cache::builder()
98 .max_capacity(16384)
99 .expire_after(DataExpiry)
100 .build();
101
102 let local_task = tokio::spawn({
103 let local = local.clone();
104
105 async move {
106 loop {
107 tokio::time::sleep(Duration::from_secs(10)).await;
108 local.run_pending_tasks().await;
109 }
110 }
111 });
112
113 let local_locks = moka::future::Cache::builder().max_capacity(4096).build();
114
115 let local_locks_task = tokio::spawn({
116 let local_locks = local_locks.clone();
117
118 async move {
119 loop {
120 tokio::time::sleep(Duration::from_secs(10)).await;
121 local_locks.run_pending_tasks().await;
122 }
123 }
124 });
125
126 let local_ratelimits = moka::future::Cache::builder().max_capacity(16384).build();
127
128 let instance = Arc::new(Self {
129 client,
130 use_internal_cache: env.app_use_internal_cache,
131 local,
132 local_task,
133 local_locks,
134 local_locks_task,
135 local_ratelimits,
136 cache_calls: AtomicU64::new(0),
137 cache_latency_ns_total: AtomicU64::new(0),
138 cache_latency_ns_max: AtomicU64::new(0),
139 cache_misses: AtomicU64::new(0),
140 });
141
142 let version = instance
143 .version()
144 .await
145 .unwrap_or_else(|_| "unknown".into());
146
147 tracing::info!(
148 "cache connected (redis@{}, {}ms, moka_enabled={})",
149 version,
150 start.elapsed().as_millis(),
151 env.app_use_internal_cache
152 );
153
154 instance
155 }
156
157 pub async fn version(&self) -> Result<compact_str::CompactString, rustis::Error> {
158 let Some(client) = &self.client else {
159 return Ok("memory-only".into());
160 };
161
162 let version: String = client.info([InfoSection::Server]).await?;
163 let version = version
164 .lines()
165 .find(|line| line.starts_with("valkey_version:"))
166 .or_else(|| {
167 version
168 .lines()
169 .find(|line| line.starts_with("redis_version:"))
170 })
171 .unwrap_or("_:unknown")
172 .split_once(':')
173 .map_or("unknown", |(_, v)| v.trim())
174 .into();
175
176 Ok(version)
177 }
178
179 pub async fn ratelimit(
180 &self,
181 limit_identifier: impl AsRef<str>,
182 limit: u64,
183 limit_window: u64,
184 client: impl AsRef<str>,
185 ) -> Result<(), ApiResponse> {
186 let key = compact_str::format_compact!(
187 "ratelimit::{}::{}",
188 limit_identifier.as_ref(),
189 client.as_ref()
190 );
191
192 let now = chrono::Utc::now().timestamp();
193
194 if let Some(redis_client) = &self.client {
195 let expiry = redis_client.expiretime(&key).await.unwrap_or_default();
196 let expire_unix: u64 = if expiry > now + 2 {
197 expiry as u64
198 } else {
199 now as u64 + limit_window
200 };
201
202 let limit_used = redis_client.get::<u64>(&key).await.unwrap_or_default() + 1;
203 redis_client
204 .set_with_options(key, limit_used, None, SetExpiration::Exat(expire_unix))
205 .await?;
206
207 if limit_used >= limit {
208 let retry_after = expire_unix.saturating_sub(now as u64);
209
210 return Err(ApiResponse::error(format!(
211 "you are ratelimited, retry in {retry_after}s"
212 ))
213 .with_status(StatusCode::TOO_MANY_REQUESTS)
214 .with_header("X-RateLimit-Limit", limit.to_compact_string())
215 .with_header(
216 "X-RateLimit-Remaining",
217 limit.saturating_sub(limit_used).to_compact_string(),
218 )
219 .with_header("X-RateLimit-Reset", expire_unix.to_compact_string())
220 .with_header("Retry-After", retry_after.to_compact_string()));
221 }
222 } else {
223 let mut current_count = 0;
224 let mut expire_unix = now as u64 + limit_window;
225
226 if let Some((count, exp)) = self.local_ratelimits.get(&key).await
227 && exp > now as u64 + 2
228 {
229 current_count = count;
230 expire_unix = exp;
231 }
232
233 let limit_used = current_count + 1;
234 self.local_ratelimits
235 .insert(key, (limit_used, expire_unix))
236 .await;
237
238 if limit_used >= limit {
239 return Err(ApiResponse::error(format!(
240 "you are ratelimited, retry in {}s",
241 expire_unix.saturating_sub(now as u64)
242 ))
243 .with_status(StatusCode::TOO_MANY_REQUESTS)
244 .with_header("X-RateLimit-Limit", limit.to_compact_string())
245 .with_header(
246 "X-RateLimit-Remaining",
247 limit.saturating_sub(limit_used).to_compact_string(),
248 )
249 .with_header("X-RateLimit-Reset", expire_unix.to_compact_string())
250 .with_header(
251 "Retry-After",
252 (expire_unix.saturating_sub(now as u64)).to_compact_string(),
253 ));
254 }
255 }
256
257 Ok(())
258 }
259
260 #[tracing::instrument(skip(self))]
261 pub async fn lock(
262 &self,
263 lock_id: impl Into<compact_str::CompactString> + std::fmt::Debug,
264 ttl: Option<u64>,
265 timeout: Option<u64>,
266 ) -> Result<CacheLock, anyhow::Error> {
267 let lock_id = lock_id.into();
268 let redis_key = compact_str::format_compact!("lock::{}", lock_id);
269 let ttl_secs = ttl.unwrap_or(30);
270 let deadline = timeout.map(|ms| Instant::now() + Duration::from_millis(ms));
271
272 tracing::debug!("acquiring cache lock");
273
274 let entry = self
275 .local_locks
276 .entry(lock_id.clone())
277 .or_insert_with(async {
278 LockEntry {
279 semaphore: Arc::new(tokio::sync::Semaphore::new(1)),
280 }
281 })
282 .await
283 .into_value();
284
285 let permit = match deadline {
286 Some(dl) => {
287 let remaining = dl.saturating_duration_since(Instant::now());
288 tokio::time::timeout(remaining, entry.semaphore.acquire_owned())
289 .await
290 .map_err(|_| anyhow::anyhow!("timed out waiting for cache lock `{}`", lock_id))?
291 .map_err(|_| anyhow::anyhow!("semaphore closed for lock `{}`", lock_id))?
292 }
293 None => entry
294 .semaphore
295 .acquire_owned()
296 .await
297 .map_err(|_| anyhow::anyhow!("semaphore closed for lock `{}`", lock_id))?,
298 };
299
300 if let Some(redis_client) = &self.client {
301 match Self::try_acquire_redis_lock(redis_client, &redis_key, ttl_secs, deadline).await?
302 {
303 true => {
304 tracing::debug!("acquired redis cache lock");
305 Ok(CacheLock::new(
306 lock_id,
307 Some(redis_client.clone()),
308 permit,
309 ttl,
310 ))
311 }
312 false => anyhow::bail!("timed out acquiring redis lock `{}`", lock_id),
313 }
314 } else {
315 tracing::debug!("acquired memory cache lock");
316 Ok(CacheLock::new(lock_id, None, permit, ttl))
317 }
318 }
319
320 async fn try_acquire_redis_lock(
321 client: &Arc<Client>,
322 redis_key: &compact_str::CompactString,
323 ttl_secs: u64,
324 deadline: Option<Instant>,
325 ) -> Result<bool, anyhow::Error> {
326 loop {
327 let acquired = client
328 .set_with_options(
329 redis_key.as_str(),
330 "1",
331 SetCondition::NX,
332 SetExpiration::Ex(ttl_secs),
333 )
334 .await
335 .unwrap_or(false);
336
337 if acquired {
338 return Ok(true);
339 }
340
341 if let Some(dl) = deadline {
342 let remaining = dl.saturating_duration_since(Instant::now());
343 if remaining.is_zero() {
344 return Ok(false);
345 }
346 tokio::time::sleep(remaining.min(Duration::from_millis(50))).await;
347 } else {
348 tokio::time::sleep(Duration::from_millis(50)).await;
349 }
350 }
351 }
352
353 #[tracing::instrument(skip(self, fn_compute))]
354 pub async fn cached<
355 T: Serialize + DeserializeOwned + Send,
356 F: FnOnce() -> Fut,
357 Fut: Future<Output = Result<T, FutErr>>,
358 FutErr: Into<anyhow::Error> + Send + Sync + 'static,
359 >(
360 &self,
361 key: &str,
362 ttl: u64,
363 fn_compute: F,
364 ) -> Result<T, anyhow::Error> {
365 let effective_moka_ttl = if self.use_internal_cache {
366 Duration::from_secs(ttl)
367 } else {
368 Duration::from_millis(50)
369 };
370
371 let client_opt = self.client.clone();
372
373 self.cache_calls.fetch_add(1, Ordering::Relaxed);
374 let start_time = Instant::now();
375
376 let entry = self
377 .local
378 .try_get_with(key.to_compact_string(), async move {
379 if let Some(client) = &client_opt {
380 tracing::debug!("checking redis cache");
381 let cached_value: Option<BulkString> = client
382 .get(key)
383 .await
384 .map_err(|err| {
385 tracing::error!("redis get error: {:?}", err);
386 err
387 })
388 .ok()
389 .flatten();
390
391 if let Some(value) = cached_value {
392 tracing::debug!("found in redis cache");
393 return Ok(DataEntry {
394 data: Arc::new(value.to_vec()),
395 intended_ttl: effective_moka_ttl,
396 });
397 }
398 }
399
400 self.cache_misses.fetch_add(1, Ordering::Relaxed);
401
402 tracing::debug!("executing compute");
403 let result = fn_compute().await.map_err(|e| e.into())?;
404 tracing::debug!("executed compute");
405
406 let serialized = rmp_serde::to_vec(&result)?;
407 let serialized_arc = Arc::new(serialized);
408
409 if let Some(client) = &client_opt {
410 let _ = client
411 .set_with_options(
412 key,
413 BulkStringRef(&serialized_arc),
414 None,
415 SetExpiration::Ex(ttl),
416 )
417 .await;
418 }
419
420 Ok::<_, anyhow::Error>(DataEntry {
421 data: serialized_arc,
422 intended_ttl: effective_moka_ttl,
423 })
424 })
425 .await;
426
427 let elapsed_ns = start_time.elapsed().as_nanos() as u64;
428 self.cache_latency_ns_total
429 .fetch_add(elapsed_ns, Ordering::Relaxed);
430
431 let _ = self.cache_latency_ns_max.fetch_update(
432 Ordering::Relaxed,
433 Ordering::Relaxed,
434 |current_max| {
435 if elapsed_ns > current_max {
436 Some(elapsed_ns)
437 } else {
438 Some(current_max)
439 }
440 },
441 );
442
443 match entry {
444 Ok(internal_entry) => Ok(rmp_serde::from_slice::<T>(&internal_entry.data)?),
445 Err(arc_error) => Err(anyhow::anyhow!("cache computation failed: {:?}", arc_error)),
446 }
447 }
448
449 pub async fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, anyhow::Error> {
450 if let Some(entry) = self.local.get(key).await {
451 tracing::debug!("get: found in moka cache");
452 return Ok(Some(rmp_serde::from_slice::<T>(&entry.data)?));
453 }
454
455 if let Some(client) = &self.client {
456 tracing::debug!("get: checking redis cache");
457 let cached_value: Option<BulkString> = client.get(key).await?;
458
459 if let Some(value) = cached_value {
460 tracing::debug!("get: found in redis cache");
461 let data = Arc::new(value.to_vec());
462 return Ok(Some(rmp_serde::from_slice::<T>(&data)?));
463 }
464 }
465
466 Ok(None)
467 }
468
469 pub async fn get_raw(&self, key: &str) -> Result<Option<Arc<Vec<u8>>>, anyhow::Error> {
470 if let Some(entry) = self.local.get(key).await {
471 tracing::debug!("get_raw: found in moka cache");
472 return Ok(Some(entry.data.clone()));
473 }
474
475 if let Some(client) = &self.client {
476 tracing::debug!("get_raw: checking redis cache");
477 let cached_value: Option<BulkString> = client.get(key).await?;
478
479 if let Some(value) = cached_value {
480 tracing::debug!("get_raw: found in redis cache");
481 return Ok(Some(Arc::new(value.to_vec())));
482 }
483 }
484
485 Ok(None)
486 }
487
488 pub async fn set<T: Serialize + Send + Sync>(
489 &self,
490 key: &str,
491 ttl: u64,
492 value: &T,
493 ) -> Result<(), anyhow::Error> {
494 let serialized = rmp_serde::to_vec(value)?;
495 let serialized_arc = Arc::new(serialized);
496
497 let effective_moka_ttl = if self.use_internal_cache {
498 Duration::from_secs(ttl)
499 } else {
500 Duration::from_millis(50)
501 };
502
503 self.local
504 .insert(
505 key.to_compact_string(),
506 DataEntry {
507 data: serialized_arc.clone(),
508 intended_ttl: effective_moka_ttl,
509 },
510 )
511 .await;
512
513 if let Some(client) = &self.client {
514 client
515 .set_with_options(
516 key,
517 BulkStringRef(&serialized_arc),
518 None,
519 SetExpiration::Ex(ttl),
520 )
521 .await?;
522 }
523
524 Ok(())
525 }
526
527 pub async fn set_raw(
528 &self,
529 key: &str,
530 ttl: u64,
531 value: impl Into<Arc<Vec<u8>>>,
532 ) -> Result<(), anyhow::Error> {
533 let serialized_arc = value.into();
534
535 let effective_moka_ttl = if self.use_internal_cache {
536 Duration::from_secs(ttl)
537 } else {
538 Duration::from_millis(50)
539 };
540
541 self.local
542 .insert(
543 key.to_compact_string(),
544 DataEntry {
545 data: serialized_arc.clone(),
546 intended_ttl: effective_moka_ttl,
547 },
548 )
549 .await;
550
551 if let Some(client) = &self.client {
552 client
553 .set_with_options(
554 key,
555 BulkStringRef(&serialized_arc),
556 None,
557 SetExpiration::Ex(ttl),
558 )
559 .await?;
560 }
561
562 Ok(())
563 }
564
565 pub async fn exists(&self, key: &str) -> Result<bool, anyhow::Error> {
566 if self.local.contains_key(key) {
567 return Ok(true);
568 }
569
570 if let Some(client) = &self.client {
571 Ok(client.exists(key).await? > 0)
572 } else {
573 Ok(false)
574 }
575 }
576
577 pub async fn list(
578 &self,
579 prefix: &str,
580 ) -> Result<Vec<compact_str::CompactString>, anyhow::Error> {
581 if let Some(client) = &self.client {
582 let keys = client.keys(format!("{}*", prefix)).await?;
583 Ok(keys)
584 } else {
585 let mut keys = Vec::new();
586 for (key, _) in self.local.iter() {
587 if key.starts_with(prefix) {
588 keys.push(key.to_compact_string());
589 }
590 }
591 Ok(keys)
592 }
593 }
594
595 pub async fn invalidate(&self, key: &str) -> Result<(), anyhow::Error> {
596 self.local.invalidate(key).await;
597 if let Some(client) = &self.client {
598 client.del(key).await?;
599 }
600
601 Ok(())
602 }
603
604 #[inline]
605 pub fn cache_calls(&self) -> u64 {
606 self.cache_calls.load(Ordering::Relaxed)
607 }
608
609 #[inline]
610 pub fn cache_misses(&self) -> u64 {
611 self.cache_misses.load(Ordering::Relaxed)
612 }
613
614 #[inline]
615 pub fn cache_latency_ns_average(&self) -> u64 {
616 let calls = self.cache_calls();
617 self.cache_latency_ns_total
618 .load(Ordering::Relaxed)
619 .checked_div(calls)
620 .unwrap_or(0)
621 }
622
623 #[inline]
624 pub fn cache_latency_ns_max(&self) -> u64 {
625 self.cache_latency_ns_max.load(Ordering::Relaxed)
626 }
627}
628
629impl Drop for Cache {
630 fn drop(&mut self) {
631 self.local_task.abort();
632 self.local_locks_task.abort();
633 }
634}
635
636pub struct CacheLock {
637 lock_id: Option<compact_str::CompactString>,
638 redis_client: Option<Arc<Client>>,
639 permit: Option<tokio::sync::OwnedSemaphorePermit>,
640 ttl_guard: Option<tokio::task::JoinHandle<()>>,
641}
642
643impl CacheLock {
644 fn new(
645 lock_id: compact_str::CompactString,
646 redis_client: Option<Arc<Client>>,
647 permit: tokio::sync::OwnedSemaphorePermit,
648 ttl: Option<u64>,
649 ) -> Self {
650 let ttl_guard = ttl.and_then(|secs| {
651 let lock_id_clone = lock_id.clone();
652 redis_client.clone().map(|client| {
653 tokio::spawn(async move {
654 tokio::time::sleep(Duration::from_secs(secs)).await;
655 tracing::warn!(%lock_id_clone, "cache lock TTL expired; force-releasing");
656 let redis_key = compact_str::format_compact!("lock::{}", lock_id_clone);
657 let _ = client.del(&redis_key).await;
658 })
659 })
660 });
661
662 Self {
663 lock_id: Some(lock_id),
664 redis_client,
665 permit: Some(permit),
666 ttl_guard,
667 }
668 }
669
670 #[inline]
671 pub fn is_active(&self) -> bool {
672 self.lock_id.is_some() && self.ttl_guard.as_ref().is_none_or(|h| !h.is_finished())
673 }
674}
675
676impl Drop for CacheLock {
677 fn drop(&mut self) {
678 if let Some(ttl_guard) = self.ttl_guard.take() {
679 ttl_guard.abort();
680 }
681
682 self.permit.take();
683
684 if let Some(lock_id) = self.lock_id.take()
685 && let Some(client) = self.redis_client.take()
686 {
687 tokio::spawn(async move {
688 let redis_key = compact_str::format_compact!("lock::{}", lock_id);
689 let _ = client.del(&redis_key).await;
690 });
691 }
692 }
693}