use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use crate::{CacheManager, HttpResponse, Result};
use http_cache_semantics::CachePolicy;
use redb::{Database, Durability, ReadableDatabase, TableDefinition};
use serde::{Deserialize, Serialize};
pub(crate) const TABLE: TableDefinition<&str, &[u8]> =
TableDefinition::new("http_cache_v1");
#[cfg_attr(docsrs, doc(cfg(feature = "manager-redb")))]
#[derive(Clone)]
pub struct RedbManager {
flush: Arc<FlushState>,
}
impl std::fmt::Debug for RedbManager {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RedbManager").finish_non_exhaustive()
}
}
#[derive(Debug, Deserialize, Serialize)]
struct Store {
response: HttpResponse,
policy: CachePolicy,
}
const DEFAULT_FLUSH_INTERVAL: u64 = 64;
struct FlushState {
db: Arc<Database>,
unflushed: AtomicU64,
interval: u64,
}
impl FlushState {
fn flush(&self) -> Result<()> {
let write_txn = self.db.begin_write()?;
write_txn.commit()?;
Ok(())
}
fn record_write(&self) {
if self.unflushed.fetch_add(1, Ordering::Relaxed) + 1 >= self.interval {
let n = self.unflushed.swap(0, Ordering::Relaxed);
if n == 0 {
return;
}
if let Err(e) = self.flush() {
self.unflushed.fetch_add(n, Ordering::Relaxed);
log::warn!("redb interval flush failed: {e}");
}
}
}
}
impl Drop for FlushState {
fn drop(&mut self) {
if self.unflushed.load(Ordering::Relaxed) > 0 {
if let Err(e) = self.flush() {
log::warn!("redb flush on drop failed: {e}");
}
}
}
}
impl RedbManager {
pub fn new(path: impl AsRef<Path>) -> Result<Self> {
let db = Database::create(path)?;
Self::from_database(Arc::new(db))
}
pub fn from_database(db: Arc<Database>) -> Result<Self> {
Self::from_database_with_flush_interval(db, DEFAULT_FLUSH_INTERVAL)
}
pub fn from_database_with_flush_interval(
db: Arc<Database>,
flush_interval: u64,
) -> Result<Self> {
let write_txn = db.begin_write()?;
{
let _table = write_txn.open_table(TABLE)?;
}
write_txn.commit()?;
Ok(Self {
flush: Arc::new(FlushState {
db,
unflushed: AtomicU64::new(0),
interval: flush_interval.max(1),
}),
})
}
pub async fn clear(&self) -> Result<()> {
let write_txn = self.flush.db.begin_write()?;
write_txn.delete_table(TABLE)?;
{
let _table = write_txn.open_table(TABLE)?;
}
write_txn.commit()?;
Ok(())
}
fn read_entry(&self, cache_key: &str) -> Result<Option<Store>> {
let read_txn = self.flush.db.begin_read()?;
let table = read_txn.open_table(TABLE)?;
match table.get(cache_key)? {
Some(guard) => {
match postcard::from_bytes::<Store>(guard.value()) {
Ok(store) => Ok(Some(store)),
Err(e) => {
log::debug!(
"Failed to deserialize cache entry for key \
'{cache_key}': {e}"
);
Ok(None)
}
}
}
None => Ok(None),
}
}
}
impl CacheManager for RedbManager {
async fn get(
&self,
cache_key: &str,
) -> Result<Option<(HttpResponse, CachePolicy)>> {
match self.read_entry(cache_key) {
Ok(Some(store)) => Ok(Some((store.response, store.policy))),
Ok(None) => Ok(None),
Err(e) => {
log::debug!("redb read failed for key '{cache_key}': {e}");
Ok(None)
}
}
}
async fn put(
&self,
cache_key: String,
response: HttpResponse,
policy: CachePolicy,
) -> Result<HttpResponse> {
let data = Store { response, policy };
let bytes = postcard::to_allocvec(&data)?;
let mut write_txn = self.flush.db.begin_write()?;
write_txn.set_durability(Durability::None)?;
{
let mut table = write_txn.open_table(TABLE)?;
table.insert(cache_key.as_str(), bytes.as_slice())?;
}
write_txn.commit()?;
self.flush.record_write();
Ok(data.response)
}
async fn delete(&self, cache_key: &str) -> Result<()> {
let write_txn = self.flush.db.begin_write()?;
{
let mut table = write_txn.open_table(TABLE)?;
table.remove(cache_key)?;
}
write_txn.commit()?;
Ok(())
}
}