use std::sync::Arc;
use flint_store::memory::MemoryStore;
use flint_store::ObjectStore;
use super::state::SyncerState;
use super::{LeanConfig, Syncer};
const PREFIX: &str = "tenant/gate";
const MIB: u64 = 1 << 20;
fn cfg_for(root: &std::path::Path) -> LeanConfig {
LeanConfig::new(PREFIX, root)
}
async fn syncer(store: &Arc<MemoryStore>, root: &std::path::Path) -> Syncer {
let cfg = cfg_for(root);
let state = SyncerState::open(cfg.state_dir()).unwrap();
Syncer {
store: store.clone() as Arc<dyn ObjectStore>,
cfg,
state,
lease: None,
noted_not_regular: Default::default(),
}
}
fn write_bytes(root: &std::path::Path, rel: &str, len: u64, fill: u8) {
let p = root.join(rel);
std::fs::create_dir_all(p.parent().unwrap()).unwrap();
let body: Vec<u8> = (0..len).map(|i| fill ^ (i as u8).wrapping_mul(31)).collect();
std::fs::write(p, body).unwrap();
}
struct Leg {
gate_bytes: u64,
part_parallelism: usize,
files: usize,
each: u64,
whole_put_max: u64,
put_delay_ms: u64,
}
async fn publish(leg: &Leg) -> (u64, Arc<MemoryStore>) {
let store = MemoryStore::new()
.with_part_parallelism(leg.part_parallelism)
.with_upload_inflight_max_bytes(leg.gate_bytes);
store.inject_put_whole_delay_ms(leg.put_delay_ms);
let store = Arc::new(store);
let gate = store.upload_gate().expect("the double was built with a gate");
let dir = tempfile::tempdir().unwrap();
let mut sc = syncer(&store, dir.path()).await;
sc.cfg.whole_put_max = leg.whole_put_max;
sc.checkout().await.unwrap();
for i in 0..leg.files {
write_bytes(dir.path(), &format!("ckpt/shard-{i:02}.bin"), leg.each, i as u8);
}
let r = sc.run_barrier().await.unwrap();
assert_eq!(r.uploaded.len(), leg.files, "every file published: {:?}", r.deferred);
assert!(r.deferred.is_empty());
assert_eq!(gate.held(), 0, "every permit released once the barrier is over");
(gate.high_water(), store)
}
#[tokio::test]
async fn compose_parts_never_exceed_the_budget_and_the_control_blows_past_it() {
let budget = 3 * MIB;
let leg = |gate_bytes| Leg {
gate_bytes,
part_parallelism: 8,
files: 12,
each: 4 * MIB,
whole_put_max: MIB, put_delay_ms: 0,
};
let (bounded, _) = publish(&leg(budget)).await;
assert!(bounded <= budget, "high-water {bounded} exceeded the {budget} budget");
assert!(bounded > 0, "nothing was charged: the gate is not on this path");
let (control, _) = publish(&leg(1 << 30)).await;
assert!(
control >= 4 * budget,
"the control's high-water ({control}) is not far above the budget ({budget}): \
the tree is too narrow for the bounded leg to have proved anything"
);
}
#[tokio::test]
async fn whole_bodies_never_exceed_the_budget_and_the_control_blows_past_it() {
let budget = 2 * MIB;
let leg = |gate_bytes| Leg {
gate_bytes,
part_parallelism: 8,
files: 24,
each: MIB / 2,
whole_put_max: 64 * MIB, put_delay_ms: 5, };
let (bounded, _) = publish(&leg(budget)).await;
assert!(bounded <= budget, "high-water {bounded} exceeded the {budget} budget");
assert!(bounded > 0, "nothing was charged: the gate is not on this path");
let (control, _) = publish(&leg(1 << 30)).await;
assert!(
control >= 4 * budget,
"the control's high-water ({control}) is not far above the budget ({budget}): \
the whole-body PUTs did not overlap, so the bounded leg proved nothing"
);
}
#[tokio::test]
async fn a_single_object_larger_than_the_budget_still_publishes() {
let budget = 2 * MIB;
let whole = Leg {
gate_bytes: budget,
part_parallelism: 8,
files: 1,
each: 6 * MIB,
whole_put_max: 64 * MIB,
put_delay_ms: 0,
};
let (hw, _) = tokio::time::timeout(std::time::Duration::from_secs(30), publish(&whole))
.await
.expect("a whole body above the budget deadlocked on the gate");
assert_eq!(hw, 6 * MIB, "the one clamped body, charged at its true size");
let parts = Leg {
gate_bytes: budget,
part_parallelism: 8,
files: 1,
each: 8 * MIB,
whole_put_max: 4 * MIB,
put_delay_ms: 0,
};
let (hw, store) = tokio::time::timeout(std::time::Duration::from_secs(30), publish(&parts))
.await
.expect("a part above the budget deadlocked on the gate");
assert_eq!(hw, 4 * MIB, "clamped parts run one at a time");
let dir = tempfile::tempdir().unwrap();
let mut reader = syncer(&store, dir.path()).await;
reader.cfg.whole_put_max = 4 * MIB;
reader.checkout().await.unwrap();
let got = std::fs::read(dir.path().join("ckpt/shard-00.bin")).unwrap();
let want: Vec<u8> = (0..8 * MIB).map(|i| 0u8 ^ (i as u8).wrapping_mul(31)).collect();
assert!(got == want, "the clamped compose did not round-trip byte-identically");
}
#[tokio::test]
async fn a_store_without_a_gate_offers_none() {
let store = MemoryStore::new().with_upload_inflight_max_bytes(0);
assert!(store.upload_gate().is_none());
let store = Arc::new(store);
let dir = tempfile::tempdir().unwrap();
let mut sc = syncer(&store, dir.path()).await;
sc.cfg.whole_put_max = MIB;
sc.checkout().await.unwrap();
write_bytes(dir.path(), "ckpt/shard-00.bin", 3 * MIB, 7);
let r = sc.run_barrier().await.unwrap();
assert_eq!(r.uploaded.len(), 1);
}