use std::env;
use std::fs::{File, OpenOptions, remove_file};
use std::io::{Seek as _, SeekFrom, Write as _};
use std::path::PathBuf;
use std::process;
use std::sync::mpsc;
use std::sync::{Arc, Barrier};
use std::thread;
use std::time::{Duration, Instant};
use anyhow::{Context as _, Result, anyhow, ensure};
use super::expert_store::{ExpertKey, ExpertLease, ExpertStore};
use super::gguf::ByteRange;
const PART_COUNT: usize = 6;
const PART_BYTES: usize = 16;
const SLOT_BYTES: usize = PART_COUNT * PART_BYTES;
const FILE_BYTES: usize = 2 * 4_096;
const STORE_BUDGET: usize = 16 * 1024 * 1024;
const THREAD_COUNT: usize = 8;
const COMPLETION_WAIT: Duration = Duration::from_secs(10);
const RUN_WAIT: Duration = Duration::from_secs(30);
const POLL_WAIT: Duration = Duration::from_millis(1);
struct Fixture {
path: PathBuf,
writer: Option<File>,
bytes: Vec<u8>,
}
impl Fixture {
fn new() -> Result<Self> {
let mut created = None;
for nonce in 0..1_000_u32 {
let path =
env::temp_dir().join(format!("goosedump-expert-cache-{}-{nonce}", process::id()));
match OpenOptions::new()
.read(true)
.write(true)
.create_new(true)
.open(&path)
{
Ok(file) => {
created = Some((path, file));
break;
}
Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {}
Err(error) => return Err(error).context("create expert-cache fixture"),
}
}
let (path, mut writer) = created.context("find an unused expert-cache fixture path")?;
let bytes = (0..FILE_BYTES)
.map(|index| u8::try_from((index * 17 + 3) % 251))
.collect::<Result<Vec<_>, _>>()?;
writer
.write_all(&bytes)
.context("write expert-cache fixture")?;
writer.sync_all().context("sync expert-cache fixture")?;
Ok(Self {
path,
writer: Some(writer),
bytes,
})
}
fn source(&self) -> Result<File> {
File::open(&self.path).context("open expert-cache fixture")
}
fn ranges(base: usize, limit: usize) -> Result<[ByteRange; PART_COUNT]> {
Ok([
part_range(base, 0, limit)?,
part_range(base, 1, limit)?,
part_range(base, 2, limit)?,
part_range(base, 3, limit)?,
part_range(base, 4, limit)?,
part_range(base, 5, limit)?,
])
}
fn expected(&self, range: ByteRange) -> Result<&[u8]> {
self.bytes
.get(range.start()..range.end())
.context("expert-cache fixture range is invalid")
}
fn overwrite(&mut self, ranges: [ByteRange; PART_COUNT], seed: u8) -> Result<()> {
let writer = self.writer.as_mut().context("fixture writer is closed")?;
for (index, range) in ranges.into_iter().enumerate() {
let value = seed.wrapping_add(u8::try_from(index)?);
self.bytes[range.start()..range.end()].fill(value);
writer
.seek(SeekFrom::Start(u64::try_from(range.start())?))
.context("seek expert-cache fixture")?;
writer
.write_all(&self.bytes[range.start()..range.end()])
.context("update expert-cache fixture")?;
}
writer
.sync_all()
.context("sync expert-cache fixture update")
}
}
impl Drop for Fixture {
fn drop(&mut self) {
self.writer.take();
let _result = remove_file(&self.path);
}
}
fn part_range(base: usize, index: usize, limit: usize) -> Result<ByteRange> {
let offset = index
.checked_mul(PART_BYTES)
.and_then(|offset| base.checked_add(offset))
.context("expert-cache fixture range overflow")?;
ByteRange::new(offset, PART_BYTES, limit)
}
fn new_store(fixture: &Fixture, max_slots: usize) -> Result<ExpertStore> {
ExpertStore::new(fixture.source()?, SLOT_BYTES, STORE_BUDGET, 1, max_slots)
}
fn load(
store: &ExpertStore,
key: ExpertKey,
ranges: [ByteRange; PART_COUNT],
) -> Result<ExpertLease> {
store
.request_many([(key, ranges)])?
.into_iter()
.next()
.context("expert request is missing")?
.wait()
}
fn load_parts(
store: &ExpertStore,
key: ExpertKey,
ranges: [ByteRange; PART_COUNT],
) -> Result<Vec<Vec<u8>>> {
let lease = load(store, key, ranges)?;
(0..PART_COUNT)
.map(|index| lease.part(index).map(<[u8]>::to_vec))
.collect()
}
fn verify_lease(
fixture: &Fixture,
lease: &ExpertLease,
ranges: [ByteRange; PART_COUNT],
) -> Result<()> {
for (index, range) in ranges.into_iter().enumerate() {
ensure!(
lease.part(index)? == fixture.expected(range)?,
"expert-cache part {index} differs"
);
}
ensure!(
lease.part(PART_COUNT).is_err(),
"out-of-range part succeeded"
);
Ok(())
}
fn verify_parts(
fixture: &Fixture,
parts: &[Vec<u8>],
ranges: [ByteRange; PART_COUNT],
) -> Result<()> {
ensure!(parts.len() == PART_COUNT, "expert-cache part count differs");
for (part, range) in parts.iter().zip(ranges) {
ensure!(
part.as_slice() == fixture.expected(range)?,
"expert-cache bytes differ"
);
}
Ok(())
}
fn check_range_math() -> Result<()> {
let parent = ByteRange::new(128, 96, 224)?;
for index in 0..3 {
let range = parent.split(3, index)?;
ensure!(range.start() == 128 + index * 32, "partition start differs");
ensure!(range.len() == 32, "partition length differs");
}
ensure!(ByteRange::new(0, 0, 1).is_err(), "empty range succeeded");
ensure!(
ByteRange::new(usize::MAX - 1, 2, usize::MAX).is_err(),
"overflowing range succeeded"
);
ensure!(
ByteRange::new(8, 8, 15).is_err(),
"out-of-bounds range succeeded"
);
ensure!(parent.split(0, 0).is_err(), "zero-way split succeeded");
ensure!(parent.split(3, 3).is_err(), "out-of-range split succeeded");
ensure!(
ByteRange::new(0, 100, 100)?.split(3, 0).is_err(),
"uneven split succeeded"
);
Ok(())
}
fn check_construction_and_load_errors(fixture: &Fixture) -> Result<()> {
ensure!(
ExpertStore::new(fixture.source()?, 0, STORE_BUDGET, 1, 1).is_err(),
"empty cache slot succeeded"
);
ensure!(
ExpertStore::new(fixture.source()?, SLOT_BYTES, STORE_BUDGET, 0, 1).is_err(),
"zero minimum slot count succeeded"
);
ensure!(
ExpertStore::new(fixture.source()?, SLOT_BYTES, STORE_BUDGET, 2, 1).is_err(),
"inverted slot bounds succeeded"
);
ensure!(
ExpertStore::new(fixture.source()?, SLOT_BYTES, 1, 1, 1).is_err(),
"undersized cache budget succeeded"
);
let valid = Fixture::ranges(0, FILE_BYTES)?;
let invalid = Fixture::ranges(FILE_BYTES, FILE_BYTES + SLOT_BYTES)?;
let store = new_store(fixture, 1)?;
let key = ExpertKey::new(0, 0);
ensure!(
load(&store, key, invalid).is_err(),
"past-EOF load succeeded"
);
let lease = load(&store, key, valid)?;
verify_lease(fixture, &lease, valid)
}
fn check_batch_limits(fixture: &Fixture) -> Result<()> {
let first = Fixture::ranges(0, FILE_BYTES)?;
let second = Fixture::ranges(4_096, FILE_BYTES)?;
let store = new_store(fixture, 1)?;
ensure!(
store
.request_many([
(ExpertKey::new(0, 0), first),
(ExpertKey::new(0, 1), second),
])
.is_err(),
"oversized expert batch succeeded"
);
let requests =
store.request_many([(ExpertKey::new(0, 0), first), (ExpertKey::new(0, 0), first)])?;
for request in requests {
let lease = request.wait()?;
verify_lease(fixture, &lease, first)?;
}
Ok(())
}
fn check_concurrent_single_flight(fixture: &Fixture) -> Result<()> {
let ranges = Fixture::ranges(0, FILE_BYTES)?;
let expected = (0..PART_COUNT)
.map(|index| fixture.expected(ranges[index]).map(<[u8]>::to_vec))
.collect::<Result<Vec<_>>>()?;
let store = Arc::new(new_store(fixture, 2)?);
let barrier = Arc::new(Barrier::new(THREAD_COUNT));
let queued_before = store.queued_load_count();
let (completed_tx, completed_rx) = mpsc::channel();
let handles = (0..THREAD_COUNT)
.map(|_| {
let store = Arc::clone(&store);
let barrier = Arc::clone(&barrier);
let completed_tx = completed_tx.clone();
thread::spawn(move || {
barrier.wait();
let outcome = load_parts(&store, ExpertKey::new(1, 7), ranges)
.map_err(|error| error.to_string());
let _result = completed_tx.send(outcome);
})
})
.collect::<Vec<_>>();
drop(completed_tx);
for _ in 0..THREAD_COUNT {
let parts = completed_rx
.recv_timeout(COMPLETION_WAIT)
.context("wait for concurrent expert request")?
.map_err(anyhow::Error::msg)?;
ensure!(parts == expected, "concurrent expert bytes differ");
}
for handle in handles {
handle
.join()
.map_err(|_| anyhow!("concurrent expert request panicked"))?;
}
let queued_after = queued_before
.checked_add(1)
.context("expert load count overflow")?;
ensure!(
store.queued_load_count() == queued_after,
"concurrent requests did not share one expert load"
);
Ok(())
}
fn check_pinned_eviction(fixture: &mut Fixture) -> Result<()> {
let first = Fixture::ranges(0, FILE_BYTES)?;
let second = Fixture::ranges(4_096, FILE_BYTES)?;
let first_key = ExpertKey::new(2, 0);
let second_key = ExpertKey::new(2, 1);
let store = Arc::new(new_store(fixture, 1)?);
let mut requests = store
.request_many([(first_key, first), (first_key, first)])?
.into_iter();
let first_lease = requests.next().context("first pin is missing")?.wait()?;
let second_lease = requests.next().context("second pin is missing")?.wait()?;
ensure!(requests.next().is_none(), "unexpected extra pin");
let waits_before = store.slot_wait_count();
let (finished_tx, finished_rx) = mpsc::channel();
let worker_store = Arc::clone(&store);
let worker = thread::spawn(move || {
let outcome =
load_parts(&worker_store, second_key, second).map_err(|error| error.to_string());
let _result = finished_tx.send(outcome);
});
wait_for_slot_wait(&store, waits_before)?;
verify_lease(fixture, &first_lease, first)?;
let waits_with_two_pins = store.slot_wait_count();
drop(first_lease);
wait_for_slot_wait(&store, waits_with_two_pins)?;
verify_lease(fixture, &second_lease, first)?;
drop(second_lease);
let loaded = finished_rx
.recv_timeout(COMPLETION_WAIT)
.context("wait for unpinned expert load")?
.map_err(anyhow::Error::msg)?;
verify_parts(fixture, &loaded, second)?;
worker
.join()
.map_err(|_| anyhow!("pinned-eviction worker panicked"))?;
fixture.overwrite(first, 0xa0)?;
let reloaded = load(&store, first_key, first)?;
verify_lease(fixture, &reloaded, first)
}
fn check_eviction_stress(fixture: &Fixture) -> Result<()> {
let first = Fixture::ranges(0, FILE_BYTES)?;
let second = Fixture::ranges(4_096, FILE_BYTES)?;
let store = new_store(fixture, 2)?;
for iteration in 0..64 {
let ranges = if iteration % 2 == 0 { first } else { second };
let lease = load(&store, ExpertKey::new(3, iteration % 8), ranges)?;
verify_lease(fixture, &lease, ranges)?;
}
Ok(())
}
fn wait_for_slot_wait(store: &ExpertStore, previous: usize) -> Result<()> {
let deadline = Instant::now() + COMPLETION_WAIT;
while store.slot_wait_count() <= previous {
ensure!(
Instant::now() < deadline,
"expert request did not enter the slot wait"
);
thread::sleep(POLL_WAIT);
}
Ok(())
}
fn run_checks() -> Result<()> {
check_range_math().context("expert byte-range math")?;
let mut fixture = Fixture::new()?;
check_construction_and_load_errors(&fixture).context("expert load failure recovery")?;
check_batch_limits(&fixture).context("expert batch admission")?;
check_concurrent_single_flight(&fixture).context("expert concurrent single-flight")?;
check_pinned_eviction(&mut fixture).context("expert pinned eviction")?;
check_eviction_stress(&fixture).context("expert eviction stress")?;
Ok(())
}
pub(super) fn run() -> Result<()> {
let (completed_tx, completed_rx) = mpsc::channel();
let worker = thread::Builder::new()
.name("expert-cache-checks".to_owned())
.spawn(move || {
let outcome = run_checks();
let _result = completed_tx.send(outcome);
})
.context("start expert-cache checks")?;
let outcome = completed_rx
.recv_timeout(RUN_WAIT)
.context("expert-cache checks timed out")?;
worker
.join()
.map_err(|_| anyhow!("expert-cache checks panicked"))?;
outcome
}