1use std::{
2 path::{Path, PathBuf},
3 sync::Arc,
4};
5
6use async_trait::async_trait;
7use bincode::{
8 config::standard as bincode_config,
9 serde::{decode_from_slice, encode_to_vec},
10};
11use bytes::Bytes;
12use chrono::{DateTime, Utc};
13use feoxdb::{FeoxError, FeoxStore};
14use hitbox_backend::format::{Format, JsonFormat};
15use hitbox_backend::{
16 Backend, BackendError, BackendResult, CacheKeyFormat, Compressor, DeleteStatus,
17 PassthroughCompressor,
18};
19use hitbox_core::{BackendLabel, CacheKey, CacheValue, Raw};
20use serde::{Deserialize, Serialize};
21
22use crate::FeOxDbError;
23
24#[derive(Serialize, Deserialize)]
25struct SerializableCacheValue {
26 #[serde(with = "serde_bytes")]
27 data: Vec<u8>,
28 stale: Option<DateTime<Utc>>,
29 expire: Option<DateTime<Utc>>,
30}
31
32impl From<CacheValue<Raw>> for SerializableCacheValue {
33 fn from(value: CacheValue<Raw>) -> Self {
34 Self {
35 data: value.data().to_vec(),
36 stale: value.stale(),
37 expire: value.expire(),
38 }
39 }
40}
41
42impl From<SerializableCacheValue> for CacheValue<Raw> {
43 fn from(value: SerializableCacheValue) -> Self {
44 CacheValue::new(Bytes::from(value.data), value.expire, value.stale)
45 }
46}
47
48#[derive(Clone)]
72pub struct FeOxDbBackend<S = JsonFormat, C = PassthroughCompressor>
73where
74 S: Format,
75 C: Compressor,
76{
77 store: Arc<FeoxStore>,
78 key_format: CacheKeyFormat,
79 serializer: S,
80 compressor: C,
81 label: BackendLabel,
82}
83
84impl<S, C> FeOxDbBackend<S, C>
85where
86 S: Format,
87 C: Compressor,
88{
89 pub fn flush(&self) {
97 self.store.flush();
98 }
99}
100
101impl FeOxDbBackend<JsonFormat, PassthroughCompressor> {
102 pub fn builder() -> FeOxDbBackendBuilder<JsonFormat, PassthroughCompressor> {
104 FeOxDbBackendBuilder::default()
105 }
106
107 pub fn in_memory() -> Result<Self, FeOxDbError> {
118 let store = FeoxStore::builder().enable_ttl(true).build()?;
119
120 Ok(Self {
121 store: Arc::new(store),
122 key_format: CacheKeyFormat::Bitcode,
123 serializer: JsonFormat,
124 compressor: PassthroughCompressor,
125 label: BackendLabel::new_static("feoxdb"),
126 })
127 }
128}
129
130pub struct FeOxDbBackendBuilder<S = JsonFormat, C = PassthroughCompressor>
145where
146 S: Format,
147 C: Compressor,
148{
149 path: Option<PathBuf>,
150 max_file_size: Option<u64>,
151 max_memory: Option<usize>,
152 key_format: CacheKeyFormat,
153 serializer: S,
154 compressor: C,
155 label: BackendLabel,
156}
157
158impl Default for FeOxDbBackendBuilder<JsonFormat, PassthroughCompressor> {
159 fn default() -> Self {
160 Self {
161 path: None,
162 max_file_size: None,
163 max_memory: None,
164 key_format: CacheKeyFormat::Bitcode,
165 serializer: JsonFormat,
166 compressor: PassthroughCompressor,
167 label: BackendLabel::new_static("feoxdb"),
168 }
169 }
170}
171
172impl<S, C> FeOxDbBackendBuilder<S, C>
173where
174 S: Format,
175 C: Compressor,
176{
177 pub fn path(mut self, path: impl AsRef<Path>) -> Self {
182 self.path = Some(path.as_ref().to_path_buf());
183 self
184 }
185
186 pub fn max_file_size(mut self, bytes: u64) -> Self {
193 self.max_file_size = Some(bytes);
194 self
195 }
196
197 pub fn max_memory(mut self, bytes: usize) -> Self {
207 self.max_memory = Some(bytes);
208 self
209 }
210
211 pub fn key_format(mut self, format: CacheKeyFormat) -> Self {
213 self.key_format = format;
214 self
215 }
216
217 pub fn label(mut self, label: impl Into<BackendLabel>) -> Self {
219 self.label = label.into();
220 self
221 }
222
223 pub fn value_format<NewS>(self, serializer: NewS) -> FeOxDbBackendBuilder<NewS, C>
228 where
229 NewS: Format,
230 {
231 FeOxDbBackendBuilder {
232 path: self.path,
233 max_file_size: self.max_file_size,
234 max_memory: self.max_memory,
235 key_format: self.key_format,
236 serializer,
237 compressor: self.compressor,
238 label: self.label,
239 }
240 }
241
242 pub fn compressor<NewC>(self, compressor: NewC) -> FeOxDbBackendBuilder<S, NewC>
248 where
249 NewC: Compressor,
250 {
251 FeOxDbBackendBuilder {
252 path: self.path,
253 max_file_size: self.max_file_size,
254 max_memory: self.max_memory,
255 key_format: self.key_format,
256 serializer: self.serializer,
257 compressor,
258 label: self.label,
259 }
260 }
261
262 pub fn build(self) -> Result<FeOxDbBackend<S, C>, FeOxDbError> {
266 let mut builder = FeoxStore::builder().enable_ttl(true);
267
268 if let Some(mut path) = self.path {
269 if path.is_dir() {
270 path.push("cache.db");
271 }
272 let path_str = path.to_string_lossy().to_string();
273 builder = builder.device_path(path_str);
274 }
275
276 if let Some(file_size) = self.max_file_size {
277 builder = builder.file_size(file_size);
278 }
279
280 if let Some(memory) = self.max_memory {
281 builder = builder.max_memory(memory);
282 }
283
284 let store = builder.build()?;
285
286 Ok(FeOxDbBackend {
287 store: Arc::new(store),
288 key_format: self.key_format,
289 serializer: self.serializer,
290 compressor: self.compressor,
291 label: self.label,
292 })
293 }
294}
295
296#[async_trait]
297impl<S, C> Backend for FeOxDbBackend<S, C>
298where
299 S: Format + Send + Sync,
300 C: Compressor + Send + Sync,
301{
302 async fn read(&self, key: &CacheKey) -> BackendResult<Option<CacheValue<Raw>>> {
303 let store = self.store.clone();
304
305 let key_bytes = encode_to_vec(key, bincode_config())
306 .map_err(|e| BackendError::InternalError(Box::new(e)))?;
307
308 tokio::task::spawn_blocking(move || match store.get(&key_bytes) {
309 Ok(encoded) => {
310 let (serializable, _): (SerializableCacheValue, _) =
311 decode_from_slice(&encoded, bincode_config())
312 .map_err(|e| BackendError::InternalError(Box::new(e)))?;
313
314 let cache_value: CacheValue<Raw> = serializable.into();
315
316 if let Some(expire_time) = cache_value.expire()
317 && expire_time < Utc::now()
318 {
319 return Ok(None);
320 }
321
322 Ok(Some(cache_value))
323 }
324 Err(FeoxError::KeyNotFound) => Ok(None),
325 Err(e) => Err(BackendError::InternalError(Box::new(e))),
326 })
327 .await
328 .map_err(|e| BackendError::InternalError(Box::new(e)))?
329 }
330
331 async fn write(&self, key: &CacheKey, value: CacheValue<Raw>) -> BackendResult<()> {
332 let store = self.store.clone();
333
334 let key_bytes = encode_to_vec(key, bincode_config())
335 .map_err(|e| BackendError::InternalError(Box::new(e)))?;
336
337 let ttl = value.ttl();
339
340 let serializable: SerializableCacheValue = value.into();
341 let value_bytes = encode_to_vec(&serializable, bincode_config())
342 .map_err(|e| BackendError::InternalError(Box::new(e)))?;
343
344 tokio::task::spawn_blocking(move || {
345 ttl.map(|ttl_duration| ttl_duration.as_secs())
346 .map(|ttl_secs| store.insert_with_ttl(&key_bytes, &value_bytes, ttl_secs))
347 .unwrap_or_else(|| store.insert(&key_bytes, &value_bytes))
348 .map_err(|e| BackendError::InternalError(Box::new(e)))?;
349 Ok(())
350 })
351 .await
352 .map_err(|e| BackendError::InternalError(Box::new(e)))?
353 }
354
355 async fn remove(&self, key: &CacheKey) -> BackendResult<DeleteStatus> {
356 let store = self.store.clone();
357
358 let key_bytes = encode_to_vec(key, bincode_config())
359 .map_err(|e| BackendError::InternalError(Box::new(e)))?;
360
361 tokio::task::spawn_blocking(move || match store.delete(&key_bytes) {
362 Ok(()) => Ok(DeleteStatus::Deleted(1)),
363 Err(FeoxError::KeyNotFound) => Ok(DeleteStatus::Missing),
364 Err(e) => Err(BackendError::InternalError(Box::new(e))),
365 })
366 .await
367 .map_err(|e| BackendError::InternalError(Box::new(e)))?
368 }
369
370 fn value_format(&self) -> &dyn Format {
371 &self.serializer
372 }
373
374 fn key_format(&self) -> &CacheKeyFormat {
375 &self.key_format
376 }
377
378 fn compressor(&self) -> &dyn Compressor {
379 &self.compressor
380 }
381
382 fn label(&self) -> BackendLabel {
383 self.label.clone()
384 }
385}
386
387impl<S, C> hitbox_backend::CacheBackend for FeOxDbBackend<S, C>
389where
390 S: Format + Send + Sync,
391 C: Compressor + Send + Sync,
392{
393}
394
395#[cfg(test)]
396mod tests {
397 use super::*;
398 use chrono::Utc;
399 use tempfile::TempDir;
400
401 #[tokio::test]
402 async fn test_write_and_read() {
403 let temp_dir = TempDir::new().unwrap();
404 let backend = FeOxDbBackend::builder()
405 .path(temp_dir.path())
406 .build()
407 .unwrap();
408
409 let key = CacheKey::from_str("test-key", "1");
410 let value = CacheValue::new(
411 Bytes::from(&b"test-value"[..]),
412 Some(Utc::now() + chrono::Duration::hours(1)),
413 None,
414 );
415
416 backend.write(&key, value.clone()).await.unwrap();
418
419 let result = backend.read(&key).await.unwrap();
421 assert!(result.is_some());
422 assert_eq!(result.unwrap().data().as_ref(), b"test-value");
423 }
424
425 #[tokio::test]
426 async fn test_delete() {
427 let temp_dir = TempDir::new().unwrap();
428 let backend = FeOxDbBackend::builder()
429 .path(temp_dir.path())
430 .build()
431 .unwrap();
432
433 let key = CacheKey::from_str("delete-key", "1");
434 let value = CacheValue::new(
435 Bytes::from(&b"test-value"[..]),
436 Some(Utc::now() + chrono::Duration::hours(1)),
437 None,
438 );
439
440 backend.write(&key, value).await.unwrap();
442
443 let status = backend.remove(&key).await.unwrap();
445 assert_eq!(status, DeleteStatus::Deleted(1));
446
447 let result = backend.read(&key).await.unwrap();
449 assert!(result.is_none());
450 }
451
452 #[tokio::test]
453 async fn test_delete_missing() {
454 let temp_dir = TempDir::new().unwrap();
455 let backend = FeOxDbBackend::builder()
456 .path(temp_dir.path())
457 .build()
458 .unwrap();
459
460 let key = CacheKey::from_str("nonexistent", "1");
461 let status = backend.remove(&key).await.unwrap();
462 assert_eq!(status, DeleteStatus::Missing);
463 }
464
465 #[tokio::test]
466 async fn test_read_nonexistent() {
467 let temp_dir = TempDir::new().unwrap();
468 let backend = FeOxDbBackend::builder()
469 .path(temp_dir.path())
470 .build()
471 .unwrap();
472
473 let key = CacheKey::from_str("nonexistent-read", "1");
474 let result = backend.read(&key).await.unwrap();
475 assert!(result.is_none());
476 }
477
478 #[tokio::test]
479 async fn test_in_memory_backend() {
480 let backend = FeOxDbBackend::in_memory().unwrap();
481
482 let key = CacheKey::from_str("memory-key", "1");
483 let value = CacheValue::new(
484 Bytes::from(&b"memory-value"[..]),
485 Some(Utc::now() + chrono::Duration::hours(1)),
486 None,
487 );
488
489 backend.write(&key, value).await.unwrap();
491
492 let result = backend.read(&key).await.unwrap();
494 assert!(result.is_some());
495 assert_eq!(result.unwrap().data().as_ref(), b"memory-value");
496 }
497
498 #[tokio::test]
499 async fn test_clone_shares_store() {
500 let temp_dir = TempDir::new().unwrap();
501 let backend1 = FeOxDbBackend::builder()
502 .path(temp_dir.path())
503 .build()
504 .unwrap();
505 let backend2 = backend1.clone();
506
507 let key = CacheKey::from_str("shared-key", "1");
508 let value = CacheValue::new(
509 Bytes::from(&b"shared-value"[..]),
510 Some(Utc::now() + chrono::Duration::hours(1)),
511 None,
512 );
513
514 backend1.write(&key, value).await.unwrap();
516
517 let result = backend2.read(&key).await.unwrap();
519 assert!(result.is_some());
520 assert_eq!(result.unwrap().data().as_ref(), b"shared-value");
521 }
522
523 #[tokio::test]
524 async fn test_per_key_ttl() {
525 let temp_dir = TempDir::new().unwrap();
526 let backend = FeOxDbBackend::builder()
527 .path(temp_dir.path())
528 .build()
529 .unwrap();
530
531 let now = Utc::now();
532 let expire_1h = now + chrono::Duration::hours(1);
533 let expire_24h = now + chrono::Duration::hours(24);
534
535 let key1 = CacheKey::from_str("key1", "1");
537 let value1 = CacheValue::new(Bytes::from(&b"value1"[..]), Some(expire_1h), None);
538 backend.write(&key1, value1).await.unwrap();
539
540 let key2 = CacheKey::from_str("key2", "1");
542 let value2 = CacheValue::new(Bytes::from(&b"value2"[..]), Some(expire_24h), None);
543 backend.write(&key2, value2).await.unwrap();
544
545 let read1 = backend
547 .read(&key1)
548 .await
549 .unwrap()
550 .expect("key1 should exist");
551 let read2 = backend
552 .read(&key2)
553 .await
554 .unwrap()
555 .expect("key2 should exist");
556
557 let tolerance = chrono::Duration::seconds(1);
559 assert!(
560 (read1.expire().unwrap() - expire_1h).abs() < tolerance,
561 "key1 expire time should be ~1 hour from now"
562 );
563 assert!(
564 (read2.expire().unwrap() - expire_24h).abs() < tolerance,
565 "key2 expire time should be ~24 hours from now"
566 );
567 }
568
569 #[tokio::test]
570 async fn test_expired_entry_not_returned() {
571 let backend = FeOxDbBackend::in_memory().unwrap();
572
573 let key = CacheKey::from_str("expired-key", "1");
575 let expired_time = Utc::now() - chrono::Duration::seconds(10);
576 let value = CacheValue::new(Bytes::from(&b"expired"[..]), Some(expired_time), None);
577 backend.write(&key, value).await.unwrap();
578
579 let result = backend.read(&key).await.unwrap();
581 assert!(result.is_none(), "Expired entry should not be returned");
582 }
583
584 #[tokio::test]
585 async fn test_memory_limit_exceeded() {
586 let backend = FeOxDbBackend::builder()
588 .max_memory(1024) .build()
590 .unwrap();
591
592 let key = CacheKey::from_str("big-key", "1");
594 let large_data = vec![0u8; 2048]; let value = CacheValue::new(
596 Bytes::from(large_data),
597 Some(Utc::now() + chrono::Duration::hours(1)),
598 None,
599 );
600
601 let result = backend.write(&key, value).await;
602 assert!(
603 result.is_err(),
604 "Write should fail when exceeding memory limit"
605 );
606 }
607
608 #[tokio::test]
609 async fn test_builder_with_label() {
610 let backend = FeOxDbBackend::builder()
611 .label("custom-label")
612 .build()
613 .unwrap();
614
615 assert_eq!(backend.label().as_ref(), "custom-label");
616 }
617
618 #[tokio::test]
619 async fn test_builder_with_custom_format() {
620 use hitbox_backend::format::BincodeFormat;
621
622 let temp_dir = TempDir::new().unwrap();
623 let backend = FeOxDbBackend::builder()
624 .path(temp_dir.path())
625 .value_format(BincodeFormat)
626 .build()
627 .unwrap();
628
629 let key = CacheKey::from_str("format-key", "1");
631 let value = CacheValue::new(
632 Bytes::from(&b"format-value"[..]),
633 Some(Utc::now() + chrono::Duration::hours(1)),
634 None,
635 );
636
637 backend.write(&key, value).await.unwrap();
638 let result = backend.read(&key).await.unwrap();
639 assert!(result.is_some());
640 assert_eq!(result.unwrap().data().as_ref(), b"format-value");
641 }
642
643 #[tokio::test]
644 async fn test_flush_persists_data() {
645 let temp_dir = TempDir::new().unwrap();
646 let db_path = temp_dir.path().join("cache.db");
647
648 {
650 let backend = FeOxDbBackend::builder()
651 .path(temp_dir.path())
652 .build()
653 .unwrap();
654
655 let key = CacheKey::from_str("persist-key", "1");
656 let value = CacheValue::new(
657 Bytes::from(&b"persist-value"[..]),
658 Some(Utc::now() + chrono::Duration::hours(1)),
659 None,
660 );
661 backend.write(&key, value).await.unwrap();
662 backend.flush();
663 }
664
665 let backend = FeOxDbBackend::builder().path(&db_path).build().unwrap();
667
668 let key = CacheKey::from_str("persist-key", "1");
669 let result = backend.read(&key).await.unwrap();
670 assert!(
671 result.is_some(),
672 "Data should persist after flush and reopen"
673 );
674 assert_eq!(result.unwrap().data().as_ref(), b"persist-value");
675 }
676
677 #[tokio::test]
678 async fn test_file_size_limit_drops_excess_writes() {
679 let temp_dir = TempDir::new().unwrap();
680 let db_path = temp_dir.path().join("cache.db");
681
682 let file_size_limit = 10 * 1024 * 1024; let chunk_size = 256 * 1024; let num_chunks = 60; {
688 let backend = FeOxDbBackend::builder()
689 .path(temp_dir.path())
690 .max_file_size(file_size_limit)
691 .build()
692 .unwrap();
693
694 let chunk = vec![0u8; chunk_size];
695 for i in 0..num_chunks {
696 let key = CacheKey::from_str(&format!("chunk-{}", i), "1");
697 let value = CacheValue::new(
698 Bytes::from(chunk.clone()),
699 Some(Utc::now() + chrono::Duration::hours(1)),
700 None,
701 );
702 let _ = backend.write(&key, value).await;
703 if i % 5 == 4 {
705 backend.flush();
706 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
707 }
708 }
709 backend.flush();
710 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
711 }
712
713 let backend = FeOxDbBackend::builder()
715 .path(&db_path)
716 .max_file_size(file_size_limit)
717 .build()
718 .unwrap();
719
720 let mut persisted_count = 0;
721 for i in 0..num_chunks {
722 let key = CacheKey::from_str(&format!("chunk-{}", i), "1");
723 if backend.read(&key).await.unwrap().is_some() {
724 persisted_count += 1;
725 }
726 }
727
728 assert!(persisted_count > 0, "At least some chunks should persist");
730 assert!(
731 persisted_count < num_chunks,
732 "Not all chunks should persist when exceeding file size limit. \
733 Persisted {}/{} chunks",
734 persisted_count,
735 num_chunks
736 );
737 }
738}