use std::collections::{HashMap, HashSet};
use std::io::Read;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Mutex;
use bytes::Bytes;
use super::chunk::is_chunk_address;
pub const MAX_CACHED_BODY: u64 = 64 * 1024 * 1024;
pub trait ChunkCache: Send + Sync {
fn get(&self, addr: &str) -> Option<Bytes>;
fn put(&self, addr: &str, body: &[u8]);
fn remove(&self, addr: &str);
fn retain(&self, live: &HashSet<String>);
}
#[derive(Default)]
pub struct MemoryChunkCache {
map: Mutex<HashMap<String, Bytes>>,
}
impl MemoryChunkCache {
pub fn len(&self) -> usize {
self.map.lock().unwrap().len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl ChunkCache for MemoryChunkCache {
fn get(&self, addr: &str) -> Option<Bytes> {
self.map.lock().unwrap().get(addr).cloned()
}
fn put(&self, addr: &str, body: &[u8]) {
if !is_chunk_address(addr) || body.len() as u64 > MAX_CACHED_BODY {
return;
}
self.map.lock().unwrap().insert(addr.to_string(), Bytes::copy_from_slice(body));
}
fn remove(&self, addr: &str) {
self.map.lock().unwrap().remove(addr);
}
fn retain(&self, live: &HashSet<String>) {
self.map.lock().unwrap().retain(|k, _| live.contains(k));
}
}
pub struct DiskChunkCache {
dir: PathBuf,
warned: AtomicBool,
}
impl DiskChunkCache {
pub fn new(dir: impl Into<PathBuf>) -> Self {
DiskChunkCache { dir: dir.into(), warned: AtomicBool::new(false) }
}
pub fn dir(&self) -> &std::path::Path {
&self.dir
}
fn usable(&self) -> bool {
std::fs::symlink_metadata(&self.dir).is_ok_and(|m| m.is_dir())
}
fn ensure(&self) -> bool {
if self.usable() {
return true;
}
if super::safefs::check_parent(&self.dir).is_err() {
return false;
}
match std::fs::create_dir(&self.dir) {
Ok(()) => {}
Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {}
Err(_) => return false,
}
self.usable()
}
fn warn_once(&self, what: &str) {
if !self.warned.swap(true, Ordering::Relaxed) {
eprintln!(
"flint-sync: the chunk cache at {} is not usable ({what}); manifest loads fetch \
every chunk from the bucket",
self.dir.display()
);
}
}
}
impl ChunkCache for DiskChunkCache {
fn get(&self, addr: &str) -> Option<Bytes> {
if !is_chunk_address(addr) || !self.usable() {
return None;
}
let f = super::safefs::open_beneath_nofollow(&self.dir, addr).ok()?;
let len = f.metadata().ok()?.len();
if len > MAX_CACHED_BODY {
return None;
}
let mut buf = Vec::with_capacity(len as usize);
(&f).take(MAX_CACHED_BODY + 1).read_to_end(&mut buf).ok()?;
Some(Bytes::from(buf))
}
fn put(&self, addr: &str, body: &[u8]) {
if !is_chunk_address(addr) || body.len() as u64 > MAX_CACHED_BODY {
return;
}
if !self.ensure() {
self.warn_once("its directory cannot be made");
return;
}
let path = self.dir.join(addr);
let tmp = self.dir.join(format!("{addr}.{}{}", std::process::id(), super::scan::TMP_SUFFIX));
let wrote = super::safefs::check_parent(&path)
.and_then(|_| super::safefs::write_via_tmp_fast(&path, &tmp, body, None).map(|_| ()));
if let Err(e) = wrote {
self.warn_once(&e.to_string());
}
}
fn remove(&self, addr: &str) {
if is_chunk_address(addr) && self.usable() {
let _ = std::fs::remove_file(self.dir.join(addr));
}
}
fn retain(&self, live: &HashSet<String>) {
if !self.usable() {
return;
}
let Ok(rd) = std::fs::read_dir(&self.dir) else { return };
for entry in rd.flatten() {
let name = entry.file_name();
let Some(name) = name.to_str() else { continue };
let stale = if is_chunk_address(name) {
!live.contains(name)
} else {
name.ends_with(super::scan::TMP_SUFFIX)
&& name.get(..32).is_some_and(is_chunk_address)
};
if stale {
let _ = std::fs::remove_file(entry.path());
}
}
}
}