use serde::{Deserialize, Serialize, de::DeserializeOwned};
use std::collections::HashMap;
use std::io::{BufRead, Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use xxhash_rust::xxh3::Xxh3;
use crate::report::{Call, Usage};
const VERSION: u32 = 2;
const PROBE: usize = 64;
pub fn key_of(parts: &[&[u8]]) -> u128 {
let mut h = Xxh3::new();
for p in parts {
h.update(&(p.len() as u32).to_le_bytes());
h.update(p);
}
h.digest128()
}
#[derive(Default)]
pub struct Dict {
strings: Vec<String>,
index: HashMap<String, u32>,
}
impl Dict {
pub fn from_vec(strings: Vec<String>) -> Self {
let index = strings
.iter()
.enumerate()
.map(|(i, s)| (s.clone(), i as u32))
.collect();
Self { strings, index }
}
pub fn intern(&mut self, s: &str) -> u32 {
if let Some(&i) = self.index.get(s) {
return i;
}
let i = self.strings.len() as u32;
self.strings.push(s.to_string());
self.index.insert(s.to_string(), i);
i
}
pub fn into_strings(self) -> Vec<String> {
self.strings
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CachedCall {
pub key: u128,
pub session: u32,
pub model: u32,
pub ts: Option<i64>,
pub usage: Usage,
pub estimated: bool,
}
impl CachedCall {
pub fn to_call(&self, source: &'static str, dict: &[Arc<str>]) -> Call {
Call {
source,
session: dict[self.session as usize].clone(),
model: dict[self.model as usize].clone(),
ts: self.ts.and_then(|s| chrono::DateTime::from_timestamp(s, 0)),
usage: self.usage,
estimated: self.estimated,
}
}
}
#[derive(Serialize, Deserialize)]
pub struct Entry<S> {
pub offset: u64,
pub mtime: i64,
pub boundary: Vec<u8>,
pub state: S,
pub dict: Vec<String>,
pub entries: Vec<CachedCall>,
}
#[derive(Serialize)]
struct CacheFile<'a, S: Serialize> {
version: u32,
files: &'a HashMap<PathBuf, Entry<S>>,
}
#[derive(Deserialize)]
struct CacheFileRead<S> {
version: u32,
files: HashMap<PathBuf, Entry<S>>,
}
fn cache_path(name: &str, scope: &[&Path]) -> PathBuf {
let mut h = Xxh3::new();
for p in scope {
let s = p.to_string_lossy();
h.update(&(s.len() as u32).to_le_bytes());
h.update(s.as_bytes());
}
std::env::home_dir()
.unwrap_or_else(|| PathBuf::from("~"))
.join(format!(".cache/llmstat/{name}-{:032x}.bin", h.digest128()))
}
pub fn load<S: DeserializeOwned>(name: &str, scope: &[&Path]) -> HashMap<PathBuf, Entry<S>> {
let t0 = std::time::Instant::now();
let Ok(bytes) = std::fs::read(cache_path(name, scope)) else {
return HashMap::new();
};
let n = bytes.len();
let out = match bincode::deserialize::<CacheFileRead<S>>(&bytes) {
Ok(c) if c.version == VERSION => c.files,
_ => HashMap::new(),
};
tracing::debug!(name, bytes = n, files = out.len(), elapsed = ?t0.elapsed(), "filecache load");
out
}
pub fn save<S: Serialize>(name: &str, scope: &[&Path], files: &HashMap<PathBuf, Entry<S>>) {
let file = cache_path(name, scope);
if let Some(dir) = file.parent()
&& std::fs::create_dir_all(dir).is_err()
{
return;
}
let c = CacheFile {
version: VERSION,
files,
};
let Ok(bytes) = bincode::serialize(&c) else {
return;
};
tracing::debug!(
name,
bytes = bytes.len(),
files = files.len(),
"filecache save"
);
let tmp = file.with_extension("tmp");
if std::fs::write(&tmp, bytes).is_ok() {
let _ = std::fs::rename(tmp, file);
}
}
pub enum Plan<S> {
Reuse,
Resume(u64, S),
Full,
}
pub fn plan_bytes<S>(len: u64, p: &Plan<S>) -> u64 {
match p {
Plan::Reuse => 0,
Plan::Resume(o, _) => len - o,
Plan::Full => len,
}
}
pub const BAR_MIN_BYTES: u64 = 32 << 20;
pub fn mtime(path: &Path) -> i64 {
std::fs::metadata(path)
.and_then(|m| m.modified())
.ok()
.and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
.map(|d| d.as_secs() as i64)
.unwrap_or(0)
}
pub fn plan<S: Clone>(path: &Path, len: u64, cached: Option<&Entry<S>>) -> Plan<S> {
let Some(e) = cached else { return Plan::Full };
if len == e.offset && e.mtime == mtime(path) {
return Plan::Reuse;
}
if len > e.offset && probe_ok(path, e) {
return Plan::Resume(e.offset, e.state.clone());
}
Plan::Full
}
fn probe_ok<S>(path: &Path, e: &Entry<S>) -> bool {
let n = e.boundary.len() as u64;
let Ok(mut f) = std::fs::File::open(path) else {
return false;
};
let mut buf = vec![0u8; n as usize];
f.seek(SeekFrom::Start(e.offset - n)).is_ok()
&& f.read_exact(&mut buf).is_ok()
&& buf == e.boundary
}
pub fn boundary(path: &Path, offset: u64) -> Vec<u8> {
let n = offset.min(PROBE as u64);
let Ok(mut f) = std::fs::File::open(path) else {
return Vec::new();
};
let mut buf = vec![0u8; n as usize];
if f.seek(SeekFrom::Start(offset - n)).is_err() || f.read_exact(&mut buf).is_err() {
return Vec::new();
}
buf
}
pub fn read_lines(
path: &Path,
offset: u64,
is_complete: impl Fn(&str) -> bool,
mut f: impl FnMut(&str),
) -> std::io::Result<u64> {
let mut file = std::fs::File::open(path)?;
file.seek(SeekFrom::Start(offset))?;
let mut reader = std::io::BufReader::new(file);
let mut consumed = offset;
let mut buf = String::new();
loop {
buf.clear();
let n = match reader.read_line(&mut buf) {
Ok(0) | Err(_) => break,
Ok(n) => n,
};
if buf.ends_with('\n') {
consumed += n as u64;
f(&buf);
} else {
if is_complete(&buf) {
consumed += n as u64;
f(&buf);
}
break;
}
}
Ok(consumed)
}