use std::collections::{HashMap, HashSet};
use std::sync::atomic::Ordering;
use std::sync::mpsc::Sender;
use superfluid_proto::linkw::ALL_SPACES;
use crate::runtime::EngineHost;
use crate::scheduler::SchedStats;
pub(crate) const LEGACY_RENEW_TICKS: u64 = 3;
pub(crate) fn renew_every(ttl_ticks: Option<u64>) -> Option<u64> {
match ttl_ticks {
None => Some(LEGACY_RENEW_TICKS),
Some(0) => None,
Some(ttl) => Some((ttl / 2).max(1)),
}
}
pub(crate) const MAX_PINS: usize = 256;
const ACQUIRE_TRIES: usize = 4;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum PinKey {
Session(u64),
Prefix(u64),
}
pub struct PinReq {
pub key: PinKey,
pub tokens: Vec<u32>,
pub want: u64,
pub deadline_unix_ms: u64,
pub anchor: Option<u64>,
pub reply: Option<Sender<u64>>,
}
const LOOKAHEAD: u32 = u32::MAX;
fn stageable<'a>(host: &EngineHost, tokens: &'a [u32]) -> &'a [u32] {
match host.max_stream_tokens() as usize {
0 => tokens,
max => &tokens[..tokens.len().min(max)],
}
}
pub fn session_target(stream: &[u32]) -> (Vec<u32>, u64) {
let mut tokens = Vec::with_capacity(stream.len() + 1);
tokens.extend_from_slice(stream);
tokens.push(LOOKAHEAD);
(tokens, stream.len() as u64)
}
pub(crate) fn common_seed_lengths(
host: &EngineHost,
m: &superfluid_proto::linkw::MatchResMsg,
cap: u64,
) -> Vec<u64> {
enum Serves {
UpTo(u64, u64),
Exactly(HashSet<u64>),
}
let mut spaces = Vec::new();
let exact = host.exact_seeds_only();
for s in &m.spaces {
let (kind, page) = host.space_shape(s.space_id).unwrap_or((0, 0));
if kind == superfluid_abi::space_kind::ENCODER_CACHE {
continue;
}
let lens = s.candidates.iter().map(|c| c.prefix_len);
spaces.push(if page > 0 && !exact {
let page = page as u64;
let longest = lens.map(|l| l / page * page).max().unwrap_or(0);
Serves::UpTo(page, longest)
} else {
Serves::Exactly(lens.collect())
});
}
if spaces.is_empty() {
return Vec::new();
}
fn gcd(a: u64, b: u64) -> u64 {
if b == 0 {
a
} else {
gcd(b, a % b)
}
}
let mut unit = 0u64;
let mut paged_max: Option<u64> = None;
for sp in &spaces {
if let Serves::UpTo(page, longest) = sp {
unit = if unit == 0 {
*page
} else {
unit / gcd(unit, *page) * page
};
paged_max = Some(paged_max.map_or(*longest, |m| m.min(*longest)));
}
}
let paged_max = paged_max.map(|m| m.min(cap) / unit * unit);
let mut lens: Vec<u64> = paged_max
.into_iter()
.chain(spaces.iter().flat_map(|sp| match sp {
Serves::UpTo(..) => Vec::new(),
Serves::Exactly(set) => set.iter().copied().collect(),
}))
.filter(|&l| l > 0 && l <= cap)
.filter(|&l| {
spaces.iter().all(|sp| match sp {
Serves::UpTo(page, longest) => l % page == 0 && l <= *longest,
Serves::Exactly(set) => set.contains(&l),
})
})
.collect();
lens.sort_unstable_by(|a, b| b.cmp(a));
lens.dedup();
lens
}
pub fn prefix_digest(tokens: &[u32]) -> u64 {
prefix_digest_from(0xCBF2_9CE4_8422_2325, tokens)
}
fn prefix_digest_from(seed: u64, tokens: &[u32]) -> u64 {
let mut h = seed;
for t in tokens {
h = superfluid_fingerprint::fnv1a64(&t.to_le_bytes(), h);
}
h
}
struct Lease {
handle: u64,
len: u64,
blocks: Vec<(u64, u64)>,
renewed_tick: u64,
renew_every: u64,
}
struct Pin {
tokens: Vec<u32>,
want: u64,
deadline_unix_ms: u64,
anchors: HashSet<u64>,
touched: u64,
lease: Option<Lease>,
}
pub struct Pins {
map: HashMap<PinKey, Pin>,
clock: u64,
budget_pct: u64,
fallback_pool_tokens: u64,
host_generation: u64,
lease_too_short: bool,
bytes: u64,
dirty: bool,
}
fn lease_blocks(host: &EngineHost, tokens: &[u32]) -> Vec<(u64, u64)> {
let mut out: Vec<(u64, u64)> = host
.paged_spaces()
.into_iter()
.flat_map(|(space_id, page, bytes)| {
let seed = superfluid_fingerprint::fnv1a64(&space_id.to_le_bytes(), 0);
block_chain_from(seed, tokens, page)
.into_iter()
.map(move |h| (h, bytes))
})
.collect();
for (space_id, bytes) in host.blob_spaces() {
let seed = superfluid_fingerprint::fnv1a64(&space_id.to_le_bytes(), 1);
out.push((prefix_digest_from(seed, tokens), bytes));
}
out
}
#[cfg(test)]
fn block_chain(tokens: &[u32], page: u64) -> Vec<u64> {
block_chain_from(0xCBF2_9CE4_8422_2325, tokens, page)
}
fn block_chain_from(seed: u64, tokens: &[u32], page: u64) -> Vec<u64> {
let mut out = Vec::with_capacity(tokens.len() / page.max(1) as usize);
let mut h = seed;
for block in tokens.chunks_exact(page.max(1) as usize) {
for t in block {
h = superfluid_fingerprint::fnv1a64(&t.to_le_bytes(), h);
}
out.push(h);
}
out
}
impl Pins {
pub fn new(budget_pct: u8, fallback_pool_tokens: u64) -> Pins {
Pins {
map: HashMap::new(),
clock: 0,
budget_pct: budget_pct as u64,
fallback_pool_tokens,
host_generation: 0,
lease_too_short: false,
bytes: 0,
dirty: false,
}
}
pub fn set_fallback_pool_tokens(&mut self, tokens: u64) {
self.fallback_pool_tokens = tokens;
}
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
pub fn until_next_deadline(&self) -> Option<std::time::Duration> {
let next = self.map.values().map(|p| p.deadline_unix_ms).min()?;
let now = crate::wal::now_unix_ms();
Some(std::time::Duration::from_millis(next.saturating_sub(now)))
}
pub fn any_held(&self) -> bool {
self.map.values().any(|p| p.lease.is_some())
}
pub fn enforceable(&self, host: &EngineHost) -> bool {
self.budget_pct > 0
&& !self.lease_too_short
&& host.page_size().is_some()
&& host.is_alive()
&& !host.has_unservable_space()
&& host.cache_serves_all_spaces()
}
fn pinned_bytes(&mut self) -> u64 {
if self.dirty {
let mut seen: HashMap<u64, u64> = HashMap::new();
for p in self.map.values() {
if let Some(l) = &p.lease {
seen.extend(l.blocks.iter().copied());
}
}
self.bytes = seen.values().sum();
self.dirty = false;
}
self.bytes
}
pub fn is_anchor(&self, session: u64) -> bool {
self.map.values().any(|p| p.anchors.contains(&session))
}
pub fn awaits_lease(&self, host: &EngineHost, session: u64) -> bool {
self.enforceable(host) && self.is_anchor(session)
}
pub fn held_max(&self, host: &EngineHost, keys: &[PinKey]) -> u64 {
if host.generation() != self.host_generation || !host.is_alive() {
return 0;
}
let now = crate::wal::now_unix_ms();
keys.iter()
.filter(|k| self.map.get(k).is_some_and(|p| p.deadline_unix_ms > now))
.map(|k| self.held(*k))
.max()
.unwrap_or(0)
}
fn held(&self, key: PinKey) -> u64 {
self.map
.get(&key)
.and_then(|p| p.lease.as_ref())
.map(|l| l.len)
.unwrap_or(0)
}
fn publish_gauges(&mut self, stats: &SchedStats) {
let held = self.map.values().filter(|p| p.lease.is_some()).count() as u64;
let bytes = self.pinned_bytes();
stats.pins_held.store(held, Ordering::Relaxed);
stats.pinned_bytes.store(bytes, Ordering::Relaxed);
}
fn check_generation(&mut self, host: &EngineHost) {
if host.generation() == self.host_generation && host.is_alive() {
return;
}
if host.generation() != self.host_generation {
self.lease_too_short = false;
}
self.host_generation = host.generation();
for p in self.map.values_mut() {
p.lease = None;
}
self.dirty = true;
}
fn release_lease(host: &mut EngineHost, lease: Lease) {
if host.is_alive() {
let _ = host.client().seed_release(lease.handle);
}
}
pub fn handle(&mut self, host: &mut EngineHost, stats: &SchedStats, tick: u64, req: PinReq) {
self.check_generation(host);
self.sweep(host, stats);
let held = if req.deadline_unix_ms == 0 {
if let Some(p) = self.map.remove(&req.key) {
if let Some(l) = p.lease {
Self::release_lease(host, l);
self.dirty = true;
}
}
0
} else {
self.clock += 1;
let clock = self.clock;
let pin = self.map.entry(req.key).or_insert_with(|| Pin {
tokens: Vec::new(),
want: 0,
deadline_unix_ms: 0,
anchors: HashSet::new(),
touched: 0,
lease: None,
});
pin.deadline_unix_ms = match req.key {
PinKey::Prefix(_) => pin.deadline_unix_ms.max(req.deadline_unix_ms),
PinKey::Session(_) => req.deadline_unix_ms,
};
pin.touched = clock;
if let Some(a) = req.anchor {
pin.anchors.insert(a);
}
if !req.tokens.is_empty() && req.want >= pin.want {
pin.tokens = req.tokens;
pin.want = req.want;
}
self.cap_entries(host, stats, req.key);
self.acquire(host, tick, req.key);
self.enforce_budget(host, stats, Some(req.key));
self.held(req.key)
};
self.publish_gauges(stats);
if let Some(reply) = req.reply {
let _ = reply.send(held);
}
}
pub fn unanchor(&mut self, session: u64) {
for (k, p) in self.map.iter_mut() {
if matches!(k, PinKey::Prefix(_)) {
p.anchors.remove(&session);
}
}
}
pub fn sweep(&mut self, host: &mut EngineHost, stats: &SchedStats) {
if self.map.is_empty() {
return;
}
if host.generation() != self.host_generation || !host.is_alive() {
self.check_generation(host);
self.publish_gauges(stats);
}
let now = crate::wal::now_unix_ms();
let expired: Vec<PinKey> = self
.map
.iter()
.filter(|(_, p)| p.deadline_unix_ms <= now)
.map(|(k, _)| *k)
.collect();
if expired.is_empty() {
return;
}
for key in expired {
if let Some(p) = self.map.remove(&key) {
if let Some(l) = p.lease {
Self::release_lease(host, l);
self.dirty = true;
}
stats.pins_expired.fetch_add(1, Ordering::Relaxed);
}
}
self.publish_gauges(stats);
}
pub fn after_tick(
&mut self,
host: &mut EngineHost,
stats: &SchedStats,
tick: u64,
retired: &[(u64, Vec<u32>)],
) {
if self.map.is_empty() {
return;
}
self.sweep(host, stats);
for (session, stream) in retired {
let keys: Vec<PinKey> = self
.map
.iter()
.filter(|(_, p)| p.anchors.contains(session))
.map(|(k, _)| *k)
.collect();
for key in keys {
if key == PinKey::Session(*session) {
if let Some(p) = self.map.get_mut(&key) {
(p.tokens, p.want) = session_target(stream);
}
} else if let Some(p) = self.map.get_mut(&key) {
p.anchors.remove(session);
}
self.acquire(host, tick, key);
}
}
self.renew(host, tick);
self.enforce_budget(host, stats, None);
self.publish_gauges(stats);
}
fn acquire(&mut self, host: &mut EngineHost, tick: u64, key: PinKey) {
if !self.enforceable(host) {
return;
}
let Some(page) = host.page_size() else { return };
let Some(pin) = self.map.get(&key) else {
return;
};
let tokens = stageable(host, &pin.tokens).to_vec();
let n = tokens.len() as u64;
let cap = pin.want.min(n.saturating_sub(1)) / page * page;
let current = pin.lease.as_ref().map(|l| l.len).unwrap_or(0);
if cap < page || current >= cap {
return;
}
let Ok(span) = host.client().stage_prompt(&tokens) else {
return;
};
let Ok(m) = host.client().match_prefix(ALL_SPACES, span) else {
return;
};
let lens = common_seed_lengths(host, &m, cap);
for len in lens
.into_iter()
.filter(|&l| l >= page && l > current)
.take(ACQUIRE_TRIES)
{
let Ok((handle, ttl)) = host.client().seed_acquire_leased(span, len, 0) else {
continue;
};
let Some(renew_every) = renew_every(ttl) else {
self.refuse_short_leases(host, handle);
return;
};
let lease = Lease {
handle,
len,
blocks: lease_blocks(host, &tokens[..len as usize]),
renewed_tick: tick,
renew_every,
};
let old = self.map.get_mut(&key).and_then(|p| p.lease.replace(lease));
if let Some(old) = old {
Self::release_lease(host, old);
}
self.dirty = true;
return;
}
}
fn renew(&mut self, host: &mut EngineHost, tick: u64) {
let due: Vec<PinKey> = self
.map
.iter()
.filter(|(_, p)| {
p.lease
.as_ref()
.is_some_and(|l| tick.saturating_sub(l.renewed_tick) >= l.renew_every)
})
.map(|(k, _)| *k)
.collect();
for key in due {
let Some(pin) = self.map.get(&key) else {
continue;
};
let Some(len) = pin.lease.as_ref().map(|l| l.len) else {
continue;
};
let tokens = stageable(host, &pin.tokens).to_vec();
let fresh = if host.is_alive() {
host.client()
.stage_prompt(&tokens)
.ok()
.and_then(|span| host.client().seed_acquire_leased(span, len, 0).ok())
} else {
None
};
let Some(pin) = self.map.get_mut(&key) else {
continue;
};
match fresh {
Some((handle, ttl)) => {
let Some(every) = renew_every(ttl) else {
if let Some(l) = pin.lease.take() {
Self::release_lease(host, l);
}
self.refuse_short_leases(host, handle);
continue;
};
let lease = pin.lease.as_mut().expect("checked above");
let old = std::mem::replace(&mut lease.handle, handle);
lease.renewed_tick = tick;
lease.renew_every = every;
let _ = host.client().seed_release(old);
}
None => {
if let Some(l) = pin.lease.take() {
Self::release_lease(host, l);
self.dirty = true;
}
}
}
}
}
fn refuse_short_leases(&mut self, host: &mut EngineHost, handle: u64) {
let _ = host.client().seed_release(handle);
for p in self.map.values_mut() {
if let Some(l) = p.lease.take() {
Self::release_lease(host, l);
}
}
self.dirty = true;
if !self.lease_too_short {
tracing::warn!(
"engine seed leases live 0 ticks: pins cannot be held across ticks and are left \
unenforced (recorded only) under this worker"
);
}
self.lease_too_short = true;
}
fn yield_oldest(
&mut self,
host: &mut EngineHost,
stats: &SchedStats,
keep: Option<PinKey>,
) -> bool {
let victim = self
.map
.iter()
.filter(|(_, p)| p.lease.is_some())
.min_by_key(|(k, p)| (Some(**k) == keep, p.touched))
.map(|(k, _)| *k);
let Some(key) = victim else { return false };
if let Some(l) = self.map.get_mut(&key).and_then(|p| p.lease.take()) {
Self::release_lease(host, l);
}
self.dirty = true;
stats.pins_yielded.fetch_add(1, Ordering::Relaxed);
true
}
fn enforce_budget(&mut self, host: &mut EngineHost, stats: &SchedStats, keep: Option<PinKey>) {
let Some(block_bytes) = host.kv_block_bytes() else {
return;
};
let mut total = stats.pool_blocks_total.load(Ordering::Relaxed);
if total == 0 {
total = self.fallback_pool_tokens / host.page_size().unwrap_or(1).max(1);
}
if total == 0 {
return;
}
let budget = total * block_bytes * self.budget_pct / 100;
while self.pinned_bytes() > budget {
if !self.yield_oldest(host, stats, keep) {
break;
}
}
}
fn cap_entries(&mut self, host: &mut EngineHost, stats: &SchedStats, keep: PinKey) {
let is_prefix = |k: &PinKey| matches!(k, PinKey::Prefix(_));
while self.map.keys().filter(|k| is_prefix(k)).count() > MAX_PINS {
let victim = self
.map
.iter()
.filter(|(k, _)| is_prefix(k) && **k != keep)
.min_by_key(|(_, p)| p.touched)
.map(|(k, _)| *k);
let Some(key) = victim else { return };
if let Some(p) = self.map.remove(&key) {
if let Some(l) = p.lease {
Self::release_lease(host, l);
self.dirty = true;
stats.pins_yielded.fetch_add(1, Ordering::Relaxed);
}
}
}
}
pub fn yield_for(&mut self, host: &mut EngineHost, stats: &SchedStats, bytes: u64) -> u64 {
let before = self.pinned_bytes();
while before - self.pinned_bytes() < bytes {
if !self.yield_oldest(host, stats, None) {
break;
}
}
self.publish_gauges(stats);
before - self.pinned_bytes()
}
pub fn yield_all(&mut self, host: &mut EngineHost, stats: &SchedStats) -> bool {
if !self.any_held() {
return false;
}
while self.yield_oldest(host, stats, None) {}
self.publish_gauges(stats);
true
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_head_only_runtime_is_seeded_only_at_an_entry_length() {
use superfluid_executor::fake::FakeConfig;
use superfluid_executor::{Executor, ExecutorConfig, FakePrimitives};
let host = |truncate_partial: bool| {
EngineHost::spawn(move || {
let cfg = FakeConfig { truncate_partial, page: 16, ..Default::default() };
(Executor::new(FakePrimitives::new(cfg), ExecutorConfig::default()), None)
})
.expect("spawn")
};
let cached_40 = superfluid_proto::linkw::MatchResMsg {
spaces: vec![superfluid_proto::linkw::SpaceMatchMsg {
space_id: 1,
candidates: vec![superfluid_proto::linkw::MatchCandidateMsg {
prefix_len: 40,
provenance_digest: [0; 32],
taint_bits: 0,
resident_tier: 0,
}],
}],
};
let paged = host(true);
assert!(!paged.exact_seeds_only());
assert_eq!(common_seed_lengths(&paged, &cached_40, 64), [32], "any page boundary within the match");
let head_only = host(false);
assert!(head_only.exact_seeds_only());
assert_eq!(common_seed_lengths(&head_only, &cached_40, 64), [40], "the entry's own length");
assert_eq!(common_seed_lengths(&head_only, &cached_40, 32), Vec::<u64>::new(), "no cut below the head");
}
#[test]
fn nested_prefixes_share_their_leading_blocks() {
let long: Vec<u32> = (0..64).collect();
let a = block_chain(&long[..32], 16);
let b = block_chain(&long, 16);
assert_eq!(a.len(), 2);
assert_eq!(b.len(), 4);
assert_eq!(
a[..],
b[..2],
"a prefix's blocks are the longer pin's leading blocks"
);
let other: Vec<u32> = (1..33).collect();
assert_ne!(block_chain(&other, 16)[0], a[0]);
assert_eq!(block_chain(&long[..40], 16).len(), 2);
}
#[test]
fn renewal_cadence_follows_the_granted_lease() {
assert_eq!(renew_every(Some(64)), Some(32));
assert_eq!(renew_every(Some(8)), Some(4));
assert_eq!(renew_every(Some(3)), Some(1));
assert_eq!(renew_every(Some(1)), Some(1));
assert_eq!(renew_every(Some(0)), None);
assert_eq!(renew_every(None), Some(LEGACY_RENEW_TICKS));
for ttl in 1..200u64 {
assert!(renew_every(Some(ttl)).unwrap() < ttl + 1);
}
}
#[test]
fn prefix_digest_is_content_identity() {
assert_eq!(prefix_digest(&[1, 2, 3]), prefix_digest(&[1, 2, 3]));
assert_ne!(prefix_digest(&[1, 2, 3]), prefix_digest(&[1, 2]));
}
}