use std::collections::HashMap;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;
use redb::{Database, Error, TableDefinition};
pub struct CacheMetadataItem {
data: Option<Vec<u8>>, media_type: Option<String>, charset: Option<String>, }
pub struct Cache {
min_file_size: usize, metadata: HashMap<String, CacheMetadataItem>, db: Option<Database>, db_ok: Option<bool>, db_file_path: Option<String>, }
const FILE_WRITE_BUF_LEN: usize = 1024 * 100; const TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("_");
impl Cache {
pub fn new(min_file_size: usize, db_file_path: Option<String>) -> Cache {
let mut cache = Cache {
min_file_size,
metadata: HashMap::new(),
db: None,
db_ok: None,
db_file_path: db_file_path.clone(),
};
if db_file_path.is_some() {
match Database::create(Path::new(&db_file_path.unwrap())) {
Ok(db) => {
cache.db = Some(db);
cache.db_ok = Some(true);
cache
}
Err(..) => {
cache.db_ok = Some(false);
cache
}
}
} else {
cache.db_ok = Some(false);
cache
}
}
pub fn set(&mut self, key: &str, data: &Vec<u8>, media_type: String, charset: String) {
let mut cache_metadata_item: CacheMetadataItem = CacheMetadataItem {
data: if self.db_ok.is_some() && self.db_ok.unwrap() {
None
} else {
Some(data.to_owned().to_vec())
},
media_type: Some(media_type.to_owned()),
charset: Some(charset),
};
if (self.db_ok.is_none() || !self.db_ok.unwrap()) || data.len() <= self.min_file_size {
cache_metadata_item.data = Some(data.to_owned().to_vec());
} else {
match self.db.as_ref().unwrap().begin_write() {
Ok(write_txn) => {
{
let mut table = write_txn.open_table(TABLE).unwrap();
table.insert(key, &*data.to_owned()).unwrap();
}
write_txn.commit().unwrap();
}
Err(..) => {
cache_metadata_item.data = Some(data.to_owned().to_vec());
}
}
}
self.metadata
.insert((*key).to_string(), cache_metadata_item);
}
pub fn get(&self, key: &str) -> Result<(Vec<u8>, String, String), Error> {
if self.metadata.contains_key(key) {
let metadata_item = self.metadata.get(key).unwrap();
if metadata_item.data.is_some() {
return Ok((
metadata_item.data.as_ref().unwrap().to_vec(),
metadata_item.media_type.as_ref().expect("").to_string(),
metadata_item.charset.as_ref().expect("").to_string(),
));
} else if self.db_ok.is_some() && self.db_ok.unwrap() {
let read_txn = self.db.as_ref().unwrap().begin_read()?;
let table = read_txn.open_table(TABLE)?;
let data = table.get(key)?;
let bytes = data.unwrap();
return Ok((
bytes.value().to_vec(),
metadata_item.media_type.as_ref().expect("").to_string(),
metadata_item.charset.as_ref().expect("").to_string(),
));
}
}
Err(Error::TransactionInProgress) }
pub fn contains_key(&self, key: &str) -> bool {
self.metadata.contains_key(key)
}
pub fn destroy_database_file(&mut self) {
if self.db_ok.is_none() || !self.db_ok.unwrap() {
return;
}
self.db = None;
self.db_ok = Some(false);
if let Some(db_file_path) = self.db_file_path.to_owned() {
if let Ok(temp_file) = File::options()
.read(true)
.write(true)
.open(db_file_path.clone())
{
let mut buffer = [0; FILE_WRITE_BUF_LEN];
let mut remaining_size: usize = temp_file.metadata().unwrap().len() as usize;
let mut writer = BufWriter::new(temp_file);
while remaining_size > 0 {
let bytes_to_write: usize = if remaining_size < FILE_WRITE_BUF_LEN {
remaining_size
} else {
FILE_WRITE_BUF_LEN
};
let buffer = &mut buffer[..bytes_to_write];
writer.write(buffer).unwrap();
remaining_size -= bytes_to_write;
}
}
}
}
}