use std::collections::{BTreeSet, HashMap, HashSet};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use jiff::SignedDuration;
use tokio::sync::watch;
use tokio::task::JoinSet;
use tracing::Instrument as _;
use tollgate_admission::{
PublicationError, PublishableSnapshotUpdate, RefreshBatch, Refreshed, SnapshotMap, Watermark,
accept_positive, accept_revoked, accept_unknown,
};
use tollgate_core::{Generation, Locality, Principal};
use tollgate_store::{Clock, SnapshotResolution, SnapshotSource, StoreError};
pub use crate::registry::SlotRegistry;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TrackedPrincipals {
Fixed(Vec<Principal>),
All {
seed: Vec<Principal>,
},
}
impl TrackedPrincipals {
fn initial(&self) -> &[Principal] {
match self {
TrackedPrincipals::Fixed(principals) => principals,
TrackedPrincipals::All { seed } => seed,
}
}
fn discovers(&self) -> bool {
matches!(self, TrackedPrincipals::All { .. })
}
}
#[derive(Debug, Clone)]
pub struct SnapshotManagerConfig {
pub principals: TrackedPrincipals,
pub refresh_interval: std::time::Duration,
pub unknown_ttl: SignedDuration,
pub revoked_ttl: SignedDuration,
pub retry_backoff: std::time::Duration,
pub max_concurrent_fetches: usize,
pub fetch_timeout: std::time::Duration,
pub enumeration_timeout: std::time::Duration,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SnapshotManagerConfigError(pub &'static str);
impl std::fmt::Display for SnapshotManagerConfigError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.0)
}
}
impl std::error::Error for SnapshotManagerConfigError {}
impl SnapshotManagerConfig {
pub fn validate(&self) -> Result<(), SnapshotManagerConfigError> {
if self.refresh_interval.is_zero() {
return Err(SnapshotManagerConfigError(
"refresh_interval must be positive",
));
}
if self.unknown_ttl <= SignedDuration::ZERO {
return Err(SnapshotManagerConfigError("unknown_ttl must be positive"));
}
if self.revoked_ttl <= SignedDuration::ZERO {
return Err(SnapshotManagerConfigError("revoked_ttl must be positive"));
}
if self.retry_backoff.is_zero() {
return Err(SnapshotManagerConfigError("retry_backoff must be positive"));
}
if self.fetch_timeout.is_zero() {
return Err(SnapshotManagerConfigError("fetch_timeout must be positive"));
}
if self.enumeration_timeout.is_zero() {
return Err(SnapshotManagerConfigError(
"enumeration_timeout must be positive",
));
}
if self.max_concurrent_fetches == 0 {
return Err(SnapshotManagerConfigError(
"max_concurrent_fetches must be positive",
));
}
let initial = self.principals.initial();
let distinct: HashSet<_> = initial.iter().copied().collect();
if distinct.len() != initial.len() {
return Err(SnapshotManagerConfigError(
"principals must not contain duplicates",
));
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SnapshotManagerReport {
pub task_died: bool,
}
#[derive(Debug)]
pub struct SnapshotCounters {
refresh_attempts: AtomicU64,
refresh_failures: AtomicU64,
refresh_timeouts: AtomicU64,
discovery_failures: AtomicU64,
refused_updates: AtomicU64,
history_evictions: AtomicU64,
publication_failures: AtomicU64,
unresolved: AtomicU64,
}
impl SnapshotCounters {
#[must_use]
pub const fn new() -> Self {
SnapshotCounters {
refresh_attempts: AtomicU64::new(0),
refresh_failures: AtomicU64::new(0),
refresh_timeouts: AtomicU64::new(0),
discovery_failures: AtomicU64::new(0),
refused_updates: AtomicU64::new(0),
history_evictions: AtomicU64::new(0),
publication_failures: AtomicU64::new(0),
unresolved: AtomicU64::new(0),
}
}
fn record_attempt(&self) {
self.refresh_attempts.fetch_add(1, Ordering::Relaxed);
}
fn record_failure(&self) {
self.refresh_failures.fetch_add(1, Ordering::Relaxed);
}
fn record_timeout(&self) {
self.refresh_timeouts.fetch_add(1, Ordering::Relaxed);
}
fn record_discovery_failure(&self) {
self.discovery_failures.fetch_add(1, Ordering::Relaxed);
}
fn set_unresolved(&self, principals: u64) {
self.unresolved.store(principals, Ordering::Relaxed);
}
#[must_use]
pub fn snapshot(&self) -> SnapshotStats {
SnapshotStats {
refresh_attempts: self.refresh_attempts.load(Ordering::Relaxed),
refresh_failures: self.refresh_failures.load(Ordering::Relaxed),
refresh_timeouts: self.refresh_timeouts.load(Ordering::Relaxed),
discovery_failures: self.discovery_failures.load(Ordering::Relaxed),
refused_updates: self.refused_updates.load(Ordering::Relaxed),
history_evictions: self.history_evictions.load(Ordering::Relaxed),
publication_failures: self.publication_failures.load(Ordering::Relaxed),
unresolved: self.unresolved.load(Ordering::Relaxed),
}
}
}
impl Default for SnapshotCounters {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SnapshotStats {
pub refresh_attempts: u64,
pub refresh_failures: u64,
pub refresh_timeouts: u64,
pub discovery_failures: u64,
pub refused_updates: u64,
pub history_evictions: u64,
pub publication_failures: u64,
pub unresolved: u64,
}
pub struct SnapshotManager {
shutdown: watch::Sender<bool>,
ready: watch::Receiver<bool>,
handle: Option<tokio::task::JoinHandle<()>>,
counters: Arc<SnapshotCounters>,
}
impl SnapshotManager {
pub fn spawn(
source: Arc<dyn SnapshotSource>,
map: Arc<dyn SnapshotMap>,
slots: Arc<SlotRegistry>,
clock: Arc<dyn Clock>,
config: SnapshotManagerConfig,
) -> Result<Self, SnapshotManagerConfigError> {
config.validate()?;
if matches!(&config.principals, TrackedPrincipals::Fixed(principals) if principals.len() > map.generation_capacity().get())
{
return Err(SnapshotManagerConfigError(
"fixed principals exceed snapshot generation capacity",
));
}
if map.local_sharding() != slots.sharding() {
return Err(SnapshotManagerConfigError(
"snapshot map and lease slots must use the same local sharding",
));
}
let (shutdown, shutdown_rx) = watch::channel(false);
let (ready_tx, ready) = crate::task_health::TaskHealth::channel(false);
let principals = config.principals.initial().len();
let counters = Arc::new(SnapshotCounters::new());
let task_counters = Arc::clone(&counters);
let handle = tokio::spawn(
run(
source,
map,
slots,
clock,
config,
shutdown_rx,
ready_tx,
task_counters,
)
.instrument(tracing::info_span!("snapshot_manager", principals)),
);
Ok(SnapshotManager {
shutdown,
ready,
handle: Some(handle),
counters,
})
}
#[must_use]
pub fn counters(&self) -> Arc<SnapshotCounters> {
Arc::clone(&self.counters)
}
#[must_use]
pub fn ready(&self) -> watch::Receiver<bool> {
self.ready.clone()
}
pub async fn shutdown(mut self) -> SnapshotManagerReport {
crate::signal(&self.shutdown, true, "snapshot-manager shutdown");
let Some(handle) = self.handle.as_mut() else {
return SnapshotManagerReport { task_died: true };
};
match handle.await {
Ok(()) => SnapshotManagerReport { task_died: false },
Err(error) => {
tracing::error!(%error, "snapshot manager task died before shutdown completed");
SnapshotManagerReport { task_died: true }
}
}
}
}
impl Drop for SnapshotManager {
fn drop(&mut self) {
if let Some(handle) = self.handle.take() {
handle.abort();
}
}
}
#[derive(Debug, Clone, Copy)]
enum Resolution {
Present {
deadline: jiff::Timestamp,
generation: Generation,
},
Negative {
deadline: jiff::Timestamp,
next_refetch: jiff::Timestamp,
watermark: Option<Watermark>,
},
}
#[derive(Clone, Copy)]
enum UpdateOrigin {
Push,
Refresh,
}
impl UpdateOrigin {
fn name(self) -> &'static str {
match self {
Self::Push => "push",
Self::Refresh => "refresh",
}
}
}
impl Resolution {
#[cfg(test)]
fn deadline(self) -> jiff::Timestamp {
match self {
Resolution::Present { deadline, .. } | Resolution::Negative { deadline, .. } => {
deadline
}
}
}
fn watermark(self) -> Option<Watermark> {
match self {
Resolution::Present { generation, .. } => Some(Watermark::Positive(generation)),
Resolution::Negative { watermark, .. } => watermark,
}
}
}
const IDLE_WAKEUP: std::time::Duration = std::time::Duration::from_secs(3_600);
enum Publications {
Push(Vec<PublishableSnapshotUpdate>),
Refreshed(Vec<Refreshed<PublishableSnapshotUpdate>>),
}
impl Publications {
fn principals(&self) -> Vec<Principal> {
match self {
Self::Push(updates) => updates
.iter()
.map(PublishableSnapshotUpdate::principal)
.collect(),
Self::Refreshed(updates) => updates.iter().map(Refreshed::principal).collect(),
}
}
}
struct Publication {
map: Arc<dyn SnapshotMap>,
slots: Arc<SlotRegistry>,
}
impl std::fmt::Debug for Publication {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Publication").finish_non_exhaustive()
}
}
impl Publication {
fn apply(&self, updates: Publications, now: jiff::Timestamp) -> Result<(), PublicationError> {
let principals = if self.slots.observes() {
updates.principals()
} else {
Vec::new()
};
match updates {
Publications::Push(updates) => self.map.apply_publishable_many_at(updates, now)?,
Publications::Refreshed(updates) => self.map.apply_refreshed_many_at(updates, now)?,
}
self.slots
.observe_many(principals.into_iter().map(|principal| {
(principal, self.map.get_at(&principal, Locality::OBSERVER))
}));
Ok(())
}
}
#[derive(Debug)]
struct Resolutions {
publication: Option<Publication>,
tracked: HashSet<Principal>,
by_principal: HashMap<Principal, Resolution>,
present: BTreeSet<(jiff::Timestamp, Principal)>,
negative: BTreeSet<(jiff::Timestamp, Principal)>,
refetch: BTreeSet<(jiff::Timestamp, Principal)>,
}
impl Resolutions {
fn new(tracked: impl IntoIterator<Item = Principal>) -> Self {
let tracked: HashSet<Principal> = tracked.into_iter().collect();
Resolutions {
publication: None,
by_principal: HashMap::new(),
tracked,
present: BTreeSet::new(),
negative: BTreeSet::new(),
refetch: BTreeSet::new(),
}
}
fn publishing(
tracked: impl IntoIterator<Item = Principal>,
map: Arc<dyn SnapshotMap>,
slots: Arc<SlotRegistry>,
) -> Self {
let mut resolutions = Self::new(tracked);
slots.retain(&resolutions.tracked);
resolutions.publication = Some(Publication { map, slots });
resolutions
}
fn publish(
&mut self,
updates: Publications,
now: jiff::Timestamp,
counters: &SnapshotCounters,
) -> bool {
let principals = updates.principals();
let result = self
.publication
.as_ref()
.expect("a running manager owns publication")
.apply(updates, now);
if let Err(error) = result {
counters
.publication_failures
.fetch_add(1, Ordering::Relaxed);
tracing::warn!(%error, "snapshot publication refused; affected principals remain unresolved and retry");
for principal in principals {
self.discard(principal);
}
false
} else {
true
}
}
fn discard(&mut self, principal: Principal) {
if let Some(previous) = self.by_principal.remove(&principal) {
self.forget(principal, previous);
}
}
fn capacity(&self) -> usize {
self.publication
.as_ref()
.expect("a running manager owns publication")
.map
.generation_capacity()
.get()
}
fn needs_refresh(&self, principal: Principal) -> bool {
!self.by_principal.contains_key(&principal)
|| self
.publication
.as_ref()
.expect("a running manager owns publication")
.map
.needs_refresh(principal)
}
fn prepare_refreshes(
&mut self,
principals: &[Principal],
counters: &SnapshotCounters,
) -> Result<RefreshBatch, PublicationError> {
let publication = self
.publication
.as_ref()
.expect("a running manager owns publication");
let batch = publication.map.prepare_refreshes(principals)?;
if !batch.evicted.is_empty() {
counters
.history_evictions
.fetch_add(batch.evicted.len() as u64, Ordering::Relaxed);
if publication.slots.observes() {
publication
.slots
.observe_many(batch.evicted.iter().map(|&principal| (principal, None)));
}
tracing::warn!(
evicted = batch.evicted.len(),
capacity = publication.map.generation_capacity().get(),
"snapshot history reclaimed; evicted principals require authoritative refresh"
);
for &principal in &batch.evicted {
self.discard(principal);
}
}
Ok(batch)
}
fn is_tracked(&self, principal: Principal) -> bool {
self.tracked.contains(&principal)
}
fn track(&mut self, principal: Principal) {
self.tracked.insert(principal);
if let Some(publication) = &self.publication {
publication.slots.track(principal);
}
}
fn retain(&mut self, discovered: HashSet<Principal>) {
let removed: Vec<Principal> = self
.tracked
.difference(&discovered)
.copied()
.collect::<Vec<_>>();
if let Some(publication) = &self.publication {
publication.map.remove_many(&removed);
if publication.slots.observes() {
publication
.slots
.observe_many(removed.iter().map(|&principal| (principal, None)));
}
}
for principal in removed {
if let Some(previous) = self.by_principal.remove(&principal) {
self.forget(principal, previous);
}
}
if let Some(publication) = &self.publication {
publication.slots.retain(&discovered);
}
self.tracked = discovered;
}
fn watermark_of(&self, principal: Principal) -> Option<Watermark> {
self.by_principal
.get(&principal)
.and_then(|resolution| resolution.watermark())
}
fn accepts_positive(
&self,
principal: Principal,
generation: Generation,
origin: UpdateOrigin,
counters: &SnapshotCounters,
) -> bool {
let current = self.by_principal.get(&principal).copied();
let (_, accepted) = accept_positive(
current.and_then(Resolution::watermark),
generation,
current.is_some_and(|resolution| matches!(resolution, Resolution::Present { .. }))
&& self
.publication
.as_ref()
.is_none_or(|publication| publication.map.contains_cached(&principal)),
);
let unchanged = matches!(current, Some(Resolution::Present { generation: held, .. })
if held == generation);
if !accepted && !unchanged {
self.record_refusal(principal, generation, "positive", origin, counters);
}
accepted
}
fn revocation(
&self,
principal: Principal,
generation: Generation,
origin: UpdateOrigin,
counters: &SnapshotCounters,
) -> (Option<Watermark>, bool) {
let decision = accept_revoked(self.watermark_of(principal), generation);
if !decision.1 {
self.record_refusal(principal, generation, "revoked", origin, counters);
}
decision
}
fn record_refusal(
&self,
principal: Principal,
generation: Generation,
kind: &'static str,
origin: UpdateOrigin,
counters: &SnapshotCounters,
) {
counters.refused_updates.fetch_add(1, Ordering::Relaxed);
tracing::warn!(
%principal,
offered = generation.0,
retained = ?self.watermark_of(principal),
origin = origin.name(),
kind,
"snapshot update refused; previous resolution retained"
);
}
fn insert(&mut self, principal: Principal, resolution: Resolution) {
if let Some(previous) = self.by_principal.insert(principal, resolution) {
self.forget(principal, previous);
}
match resolution {
Resolution::Present { deadline, .. } => {
self.present.insert((deadline, principal));
}
Resolution::Negative {
deadline,
next_refetch,
..
} => {
self.negative.insert((deadline, principal));
self.refetch.insert((next_refetch, principal));
}
}
}
fn back_off(&mut self, principal: Principal, retry_at: jiff::Timestamp) {
let Some(Resolution::Negative { next_refetch, .. }) = self.by_principal.get_mut(&principal)
else {
return;
};
let previous = std::mem::replace(next_refetch, retry_at);
self.refetch.remove(&(previous, principal));
self.refetch.insert((retry_at, principal));
}
fn forget(&mut self, principal: Principal, resolution: Resolution) {
match resolution {
Resolution::Present { deadline, .. } => {
self.present.remove(&(deadline, principal));
}
Resolution::Negative {
deadline,
next_refetch,
..
} => {
self.negative.remove(&(deadline, principal));
self.refetch.remove(&(next_refetch, principal));
}
}
}
fn expire_through(&mut self, now: jiff::Timestamp) {
for set in [&mut self.present, &mut self.negative] {
while let Some((deadline, _)) = set.first() {
if *deadline > now {
break;
}
set.pop_first();
}
}
}
fn unresolved(&mut self, now: jiff::Timestamp) -> usize {
self.expire_through(now);
self.tracked.len() - self.present.len() - self.negative.len()
}
fn next_readiness_check(&mut self, now: jiff::Timestamp) -> std::time::Duration {
self.expire_through(now);
let earliest = match (self.present.first(), self.negative.first()) {
(Some((a, _)), Some((b, _))) => Some((*a).min(*b)),
(Some((only, _)), None) | (None, Some((only, _))) => Some(*only),
(None, None) => None,
};
earliest.map_or(IDLE_WAKEUP, |deadline| until(deadline, now))
}
fn next_control_wakeup(&mut self, now: jiff::Timestamp) -> std::time::Duration {
self.expire_through(now);
let earliest = match (self.present.first(), self.refetch.first()) {
(Some((a, _)), Some((b, _))) => Some((*a).min(*b)),
(Some((only, _)), None) | (None, Some((only, _))) => Some(*only),
(None, None) => None,
};
earliest.map_or(IDLE_WAKEUP, |deadline| until(deadline, now))
}
fn due_for_sweep(&self) -> Vec<Principal> {
#[allow(
clippy::disallowed_methods,
reason = "sorted below before it is returned, so the hash order never reaches the caller"
)]
let mut principals: Vec<Principal> = self
.tracked
.iter()
.filter(|principal| {
!matches!(
self.by_principal.get(principal),
Some(Resolution::Negative { .. })
)
})
.copied()
.collect();
principals.sort_unstable();
principals
}
fn all_tracked(&self) -> Vec<Principal> {
#[allow(
clippy::disallowed_methods,
reason = "sorted on the next line before it is returned, so the hash order never reaches the caller"
)]
let mut principals: Vec<Principal> = self.tracked.iter().copied().collect();
principals.sort_unstable();
principals
}
fn due_for_refetch(&self, now: jiff::Timestamp, limit: usize) -> Vec<Principal> {
self.refetch
.range(..(next_instant(now), Principal(0)))
.take(limit)
.map(|(_, principal)| *principal)
.collect()
}
}
fn next_instant(now: jiff::Timestamp) -> jiff::Timestamp {
now.checked_add(SignedDuration::from_nanos(1))
.unwrap_or(jiff::Timestamp::MAX)
}
fn until(deadline: jiff::Timestamp, now: jiff::Timestamp) -> std::time::Duration {
if deadline <= now {
return std::time::Duration::ZERO;
}
u64::try_from(deadline.duration_since(now).as_nanos())
.map(std::time::Duration::from_nanos)
.unwrap_or(std::time::Duration::MAX)
}
fn ready_now(mode: &TrackedPrincipals, outstanding: usize, tracked: usize) -> bool {
match mode {
TrackedPrincipals::Fixed(_) => outstanding == 0,
TrackedPrincipals::All { .. } => tracked == 0 || outstanding < tracked,
}
}
fn update_ready(
ready: &watch::Sender<bool>,
resolutions: &mut Resolutions,
clock: &Arc<dyn Clock>,
counters: &SnapshotCounters,
mode: &TrackedPrincipals,
) {
let outstanding = resolutions.unresolved(clock.now());
counters.set_unresolved(outstanding as u64);
let now_ready = ready_now(mode, outstanding, resolutions.tracked.len());
if ready.borrow().ne(&now_ready) {
tracing::info!(ready = now_ready, "snapshot readiness changed");
}
crate::signal(ready, now_ready, "snapshot-manager readiness");
}
fn after_std(now: jiff::Timestamp, duration: std::time::Duration) -> jiff::Timestamp {
let nanos = i64::try_from(duration.as_nanos()).unwrap_or(i64::MAX);
now.checked_add(SignedDuration::from_nanos(nanos))
.unwrap_or(jiff::Timestamp::MAX)
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum NegativeKind {
Revoked,
Unknown,
}
fn negative_deadline(
now: jiff::Timestamp,
config: &SnapshotManagerConfig,
kind: NegativeKind,
) -> jiff::Timestamp {
let ttl = match kind {
NegativeKind::Revoked => config.revoked_ttl,
NegativeKind::Unknown => config.unknown_ttl,
};
now.checked_add(ttl).unwrap_or(jiff::Timestamp::MAX)
}
#[allow(
clippy::too_many_arguments,
reason = "one refresh pass's inputs, owned by the enclosing task; a parameter struct would rename that task's state"
)]
async fn refresh_all_cancellable(
source: &Arc<dyn SnapshotSource>,
slots: &Arc<SlotRegistry>,
clock: &Arc<dyn Clock>,
config: &SnapshotManagerConfig,
principals: &[Principal],
resolutions: &mut Resolutions,
shutdown: &mut watch::Receiver<bool>,
ready: &watch::Sender<bool>,
counters: &SnapshotCounters,
) -> Option<Vec<Principal>> {
let mut pending = Vec::new();
for principals in principals.chunks(resolutions.capacity()) {
pending.extend(
refresh_chunk_cancellable(
source,
slots,
clock,
config,
principals,
resolutions,
shutdown,
ready,
counters,
)
.await?,
);
}
Some(pending)
}
#[allow(
clippy::too_many_arguments,
reason = "each chunk takes the same set the enclosing pass owns"
)]
async fn refresh_chunk_cancellable(
source: &Arc<dyn SnapshotSource>,
slots: &Arc<SlotRegistry>,
clock: &Arc<dyn Clock>,
config: &SnapshotManagerConfig,
principals: &[Principal],
resolutions: &mut Resolutions,
shutdown: &mut watch::Receiver<bool>,
ready: &watch::Sender<bool>,
counters: &SnapshotCounters,
) -> Option<Vec<Principal>> {
let mut pending = match resolutions.prepare_refreshes(principals, counters) {
Ok(batch) => batch.reads.into_iter(),
Err(error) => {
counters
.publication_failures
.fetch_add(1, Ordering::Relaxed);
tracing::warn!(%error, "snapshot refresh reservation failed; principals retry with backoff");
return Some(principals.to_vec());
}
};
update_ready(ready, resolutions, clock, counters, &config.principals);
type Fetched = Result<Result<SnapshotResolution, StoreError>, tokio::time::error::Elapsed>;
let mut tasks = JoinSet::<Refreshed<Fetched>>::new();
for _ in 0..config.max_concurrent_fetches {
let Some(read) = pending.next() else {
break;
};
let source = Arc::clone(source);
let bound = config.fetch_timeout;
let principal = read.principal();
tasks.spawn(async move {
read.fetch(|| tokio::time::timeout(bound, source.snapshot(principal)))
.await
});
}
let mut results = Vec::with_capacity(principals.len());
while !tasks.is_empty() {
let readiness_check = tokio::time::sleep(resolutions.next_readiness_check(clock.now()));
tokio::pin!(readiness_check);
tokio::select! {
joined = tasks.join_next() => {
match joined {
Some(Ok(result)) => results.push(result),
Some(Err(error)) => {
tracing::error!(%error, "snapshot fetch task died");
}
None => {}
}
if let Some(read) = pending.next() {
let source = Arc::clone(source);
let bound = config.fetch_timeout;
let principal = read.principal();
tasks.spawn(async move {
read.fetch(|| tokio::time::timeout(bound, source.snapshot(principal))).await
});
}
}
changed = shutdown.changed() => {
if changed.is_err() || *shutdown.borrow() {
tasks.abort_all();
return None;
}
}
_ = &mut readiness_check => {
update_ready(ready, resolutions, clock, counters, &config.principals);
}
}
}
let mut updates = Vec::with_capacity(results.len());
let mut completed = HashSet::with_capacity(results.len());
for read in results {
let principal = read.principal();
let update = read.filter_map(|result| {
counters.record_attempt();
let Ok(result) = result else {
counters.record_timeout();
tracing::warn!(
%principal,
timeout_ms = config.fetch_timeout.as_millis(),
"snapshot fetch abandoned at its bound; principal keeps its \
previous resolution and retries with backoff"
);
return None;
};
match result {
Ok(SnapshotResolution::Present(snapshot)) => {
if !resolutions.accepts_positive(
principal,
snapshot.generation,
UpdateOrigin::Refresh,
counters,
) {
return None;
}
completed.insert(principal);
let slot = slots.slot(snapshot.account_id);
resolutions.insert(
principal,
Resolution::Present {
deadline: snapshot.valid_until,
generation: snapshot.generation,
},
);
Some(PublishableSnapshotUpdate::Present {
principal,
snapshot,
lease: slot,
})
}
Ok(SnapshotResolution::Revoked {
generation: incoming,
}) => {
let now = clock.now();
let (watermark, accepted) = resolutions.revocation(
principal,
incoming,
UpdateOrigin::Refresh,
counters,
);
if !accepted {
return None;
}
let until = negative_deadline(now, config, NegativeKind::Revoked);
completed.insert(principal);
resolutions.insert(
principal,
Resolution::Negative {
deadline: until,
next_refetch: until,
watermark,
},
);
Some(PublishableSnapshotUpdate::Revoked {
principal,
until,
generation: incoming,
})
}
Ok(SnapshotResolution::Unknown) => {
completed.insert(principal);
let now = clock.now();
let (watermark, _) = accept_unknown(resolutions.watermark_of(principal));
let until = negative_deadline(now, config, NegativeKind::Unknown);
resolutions.insert(
principal,
Resolution::Negative {
deadline: until,
next_refetch: until,
watermark,
},
);
Some(PublishableSnapshotUpdate::Unknown { principal, until })
}
Err(error) => {
counters.record_failure();
tracing::warn!(
%principal,
%error,
"snapshot fetch failed; principal keeps its previous resolution"
);
None
}
}
});
if let Some(update) = update {
updates.push(update);
}
}
if !updates.is_empty() {
let published = updates.iter().map(Refreshed::principal).collect::<Vec<_>>();
if !resolutions.publish(Publications::Refreshed(updates), clock.now(), counters) {
for principal in published {
completed.remove(&principal);
}
}
}
Some(
principals
.iter()
.copied()
.filter(|principal| !completed.contains(principal))
.collect(),
)
}
async fn discover(
source: &Arc<dyn SnapshotSource>,
resolutions: &mut Resolutions,
counters: &SnapshotCounters,
config: &SnapshotManagerConfig,
shutdown: &mut watch::Receiver<bool>,
) -> Option<()> {
let enumerate = || tokio::time::timeout(config.enumeration_timeout, source.principals());
let enumerated = tokio::select! {
enumerated = enumerate() => enumerated,
changed = shutdown.changed() => {
if changed.is_err() || *shutdown.borrow() {
tracing::debug!("shutdown observed during principal enumeration");
return None;
}
enumerate().await
}
};
let Ok(enumerated) = enumerated else {
counters.record_discovery_failure();
tracing::warn!(
timeout_ms = config.enumeration_timeout.as_millis(),
"principal enumeration abandoned at its bound; keeping the current set"
);
return Some(());
};
match enumerated {
Ok(Some(discovered)) => resolutions.retain(discovered.into_iter().collect()),
Ok(None) => {}
Err(error) => {
counters.record_discovery_failure();
tracing::warn!(%error, "principal enumeration failed; keeping the current set");
}
}
Some(())
}
#[allow(
clippy::too_many_arguments,
reason = "task entry point: every argument is a handle the loop owns for its lifetime, assembled once by spawn"
)]
async fn run(
source: Arc<dyn SnapshotSource>,
map: Arc<dyn SnapshotMap>,
slots: Arc<SlotRegistry>,
clock: Arc<dyn Clock>,
config: SnapshotManagerConfig,
mut shutdown: watch::Receiver<bool>,
readiness: crate::task_health::TaskHealth,
counters: Arc<SnapshotCounters>,
) {
let ready = readiness.sender();
let mut updates = source.subscribe();
let mut resolutions = Resolutions::publishing(
config.principals.initial().iter().copied(),
Arc::clone(&map),
Arc::clone(&slots),
);
if config.principals.discovers()
&& discover(&source, &mut resolutions, &counters, &config, &mut shutdown)
.await
.is_none()
{
return;
}
#[allow(
clippy::disallowed_methods,
reason = "a work queue drained to empty, not an output: every tracked principal is fetched, and the pass reports the same result whichever order they were fetched in"
)]
let mut pending: Vec<Principal> = resolutions.tracked.iter().copied().collect();
while !pending.is_empty() {
if *shutdown.borrow() {
return;
}
let Some(failed) = refresh_all_cancellable(
&source,
&slots,
&clock,
&config,
&pending,
&mut resolutions,
&mut shutdown,
ready,
&counters,
)
.await
else {
return;
};
pending = failed;
update_ready(
ready,
&mut resolutions,
&clock,
&counters,
&config.principals,
);
if !pending.is_empty() {
tokio::select! {
_ = tokio::time::sleep(config.retry_backoff) => {}
changed = shutdown.changed() => {
if changed.is_err() || *shutdown.borrow() {
return;
}
}
}
}
}
update_ready(
ready,
&mut resolutions,
&clock,
&counters,
&config.principals,
);
let mut tick = tokio::time::interval(config.refresh_interval);
tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
tick.reset(); let mut updates_closed = false;
loop {
if *shutdown.borrow() {
return;
}
let control_wakeup = tokio::time::sleep(resolutions.next_control_wakeup(clock.now()));
tokio::pin!(control_wakeup);
tokio::select! {
push = updates.recv(), if !updates_closed => match push {
Ok(push) => {
if config.principals.discovers() {
resolutions.track(push.principal);
}
if resolutions.is_tracked(push.principal) {
if resolutions.needs_refresh(push.principal) {
if refresh_all_cancellable(&source, &slots, &clock, &config, &[push.principal],
&mut resolutions, &mut shutdown, ready, &counters).await.is_none() { return; }
update_ready(ready, &mut resolutions, &clock, &counters, &config.principals);
continue;
}
match push.resolution {
SnapshotResolution::Present(snapshot) => {
if !resolutions
.accepts_positive(push.principal, snapshot.generation, UpdateOrigin::Push, &counters)
{
continue;
}
let slot = slots.slot(snapshot.account_id);
resolutions.insert(
push.principal,
Resolution::Present {
deadline: snapshot.valid_until,
generation: snapshot.generation,
},
);
resolutions.publish(
Publications::Push(vec![PublishableSnapshotUpdate::Present {
principal: push.principal,
snapshot,
lease: slot,
}]),
clock.now(),
&counters,
);
}
resolution @ (SnapshotResolution::Revoked { .. }
| SnapshotResolution::Unknown) => {
let current = resolutions.watermark_of(push.principal);
let (watermark, accepted, kind, update) = match resolution {
SnapshotResolution::Revoked { generation } => {
let (watermark, accepted) =
resolutions.revocation(push.principal, generation, UpdateOrigin::Push, &counters);
(watermark, accepted, NegativeKind::Revoked, Some(generation))
}
SnapshotResolution::Unknown => {
let (watermark, accepted) = accept_unknown(current);
(watermark, accepted, NegativeKind::Unknown, None)
}
SnapshotResolution::Present(_) => unreachable!(),
};
if !accepted {
continue;
}
let now = clock.now();
let until = negative_deadline(now, &config, kind);
resolutions.insert(
push.principal,
Resolution::Negative {
deadline: until,
next_refetch: until,
watermark,
},
);
let update = match update {
Some(generation) => PublishableSnapshotUpdate::Revoked {
principal: push.principal,
until,
generation,
},
None => PublishableSnapshotUpdate::Unknown {
principal: push.principal,
until,
},
};
resolutions.publish(Publications::Push(vec![update]), now, &counters);
}
}
update_ready(ready, &mut resolutions, &clock, &counters, &config.principals);
}
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
if refresh_all_cancellable(
&source, &slots, &clock, &config, &resolutions.all_tracked(),
&mut resolutions, &mut shutdown,
ready,
&counters,
).await.is_none() {
return;
}
update_ready(ready, &mut resolutions, &clock, &counters, &config.principals);
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
updates_closed = true;
}
},
_ = tick.tick() => {
if config.principals.discovers()
&& discover(&source, &mut resolutions, &counters, &config, &mut shutdown)
.await
.is_none()
{
return;
}
if refresh_all_cancellable(
&source, &slots, &clock, &config, &resolutions.due_for_sweep(),
&mut resolutions, &mut shutdown,
ready,
&counters,
).await.is_none() {
return;
}
update_ready(ready, &mut resolutions, &clock, &counters, &config.principals);
}
_ = &mut control_wakeup => {
let now = clock.now();
update_ready(ready, &mut resolutions, &clock, &counters, &config.principals);
let due = resolutions.due_for_refetch(now, config.max_concurrent_fetches);
if !due.is_empty() {
let Some(failed) = refresh_all_cancellable(
&source, &slots, &clock, &config, &due,
&mut resolutions, &mut shutdown, ready, &counters,
).await else {
return;
};
let retry_at = after_std(clock.now(), config.retry_backoff);
for principal in failed {
resolutions.back_off(principal, retry_at);
}
update_ready(ready, &mut resolutions, &clock, &counters, &config.principals);
}
}
changed = shutdown.changed() => {
if changed.is_err() || *shutdown.borrow() {
return;
}
}
}
}
}
#[cfg(test)]
mod tests {
use proptest::prelude::*;
use super::*;
fn t(seconds: i64) -> jiff::Timestamp {
jiff::Timestamp::from_second(seconds).unwrap()
}
#[test]
fn observing_a_publication_does_not_consume_an_affinity() {
use std::sync::Mutex;
use tollgate_admission::{
AdmissionCounters, ArcSwapSnapshotMap, LeaseSlot, MapEntry, PublicationError,
};
use tollgate_core::{
AccountId, AccountSnapshot, AccountStatus, CostTable, CostUnits, Generation,
LocalSharding, PermissionBits, PublishableSnapshot, ResolvedLimits,
};
struct RecordingMap {
inner: ArcSwapSnapshotMap,
reads: Mutex<Vec<Option<Locality>>>,
}
impl SnapshotMap for RecordingMap {
fn get(&self, principal: &Principal) -> Option<MapEntry> {
self.reads
.lock()
.expect("recording map poisoned")
.push(None);
self.inner.get(principal)
}
fn get_at(&self, principal: &Principal, locality: Locality) -> Option<MapEntry> {
self.reads
.lock()
.expect("recording map poisoned")
.push(Some(locality));
self.inner.get_at(principal, locality)
}
fn counters(&self) -> &Arc<AdmissionCounters> {
self.inner.counters()
}
fn install(
&self,
principal: Principal,
snapshot: Arc<AccountSnapshot>,
lease: Arc<LeaseSlot>,
) -> Result<(), PublicationError> {
self.inner.install(principal, snapshot, lease)
}
fn install_revoked(
&self,
principal: Principal,
until: jiff::Timestamp,
generation: Generation,
) -> Result<(), PublicationError> {
self.inner.install_revoked(principal, until, generation)
}
fn install_unknown(
&self,
principal: Principal,
until: jiff::Timestamp,
) -> Result<(), PublicationError> {
self.inner.install_unknown(principal, until)
}
fn remove(&self, principal: &Principal) {
self.inner.remove(principal);
}
}
let map = Arc::new(RecordingMap {
inner: ArcSwapSnapshotMap::new(),
reads: Mutex::new(Vec::new()),
});
let (slots, _changes) = SlotRegistry::observed(LocalSharding::SINGLE);
assert!(
slots.observes(),
"an observing registry is what reaches the read under test"
);
let publication = Publication {
map: Arc::clone(&map) as Arc<dyn SnapshotMap>,
slots,
};
let principal = Principal(7);
let snapshot = PublishableSnapshot::try_new(Arc::new(
AccountSnapshot::builder(
AccountId(1),
Generation(1),
AccountStatus::Active,
t(3_600),
PermissionBits::bit(0),
ResolvedLimits::new(16),
Arc::new(CostTable::builder(CostUnits(1), CostUnits(1)).build()),
)
.build(),
))
.expect("test snapshot limits are valid");
publication
.apply(
Publications::Push(vec![PublishableSnapshotUpdate::Present {
principal,
snapshot,
lease: LeaseSlot::for_account(AccountId(1)),
}]),
t(0),
)
.expect("the fixture publishes a well-formed snapshot");
let reads = map.reads.lock().expect("recording map poisoned").clone();
assert_eq!(
reads,
vec![Some(Locality::OBSERVER)],
"the publication must observe itself at the observer affinity, never \
through a read that resolves the calling thread's"
);
assert!(
map.get_at(&principal, Locality::OBSERVER).is_some(),
"and the publication still landed"
);
}
#[test]
fn history_reclamation_bounds_resolution_and_deadline_indexes() {
use std::num::NonZeroUsize;
use tollgate_admission::ArcSwapSnapshotMap;
use tollgate_core::LocalSharding;
let map = Arc::new(ArcSwapSnapshotMap::with_capacities(
LocalSharding::SINGLE,
7,
NonZeroUsize::new(7).unwrap(),
));
let mut resolutions =
Resolutions::publishing((0..1000).map(Principal), map.clone(), SlotRegistry::new());
let counters = SnapshotCounters::new();
for principal in (0..1000).map(Principal) {
let read = resolutions
.prepare_refreshes(&[principal], &counters)
.unwrap()
.reads
.pop()
.unwrap();
resolutions.insert(
principal,
Resolution::Negative {
deadline: t(10),
next_refetch: t(10),
watermark: None,
},
);
assert!(resolutions.publish(
Publications::Refreshed(vec![read.read(|| PublishableSnapshotUpdate::Unknown {
principal,
until: t(10)
})]),
t(0),
&counters
));
let retained = ((principal.0 + 1) as usize).min(7);
assert_eq!(resolutions.by_principal.len(), retained);
assert_eq!(resolutions.negative.len(), retained);
assert_eq!(resolutions.refetch.len(), retained);
assert!(resolutions.present.is_empty());
assert_eq!(map.history_stats().unwrap().retained, retained);
}
assert_eq!(resolutions.unresolved(t(0)), 993);
assert_eq!(counters.snapshot().history_evictions, 993);
assert_eq!(counters.snapshot().publication_failures, 0);
}
#[test]
fn a_superseded_publication_cannot_claim_a_resolved_deadline() {
use std::num::NonZeroUsize;
use tollgate_admission::ArcSwapSnapshotMap;
use tollgate_core::LocalSharding;
let map = Arc::new(ArcSwapSnapshotMap::with_capacities(
LocalSharding::SINGLE,
1,
NonZeroUsize::new(1).unwrap(),
));
let mut resolutions = Resolutions::publishing(
[Principal(1), Principal(2)],
map.clone(),
SlotRegistry::new(),
);
let counters = SnapshotCounters::new();
let stale = resolutions
.prepare_refreshes(&[Principal(1)], &counters)
.unwrap()
.reads
.pop()
.unwrap();
resolutions
.prepare_refreshes(&[Principal(2)], &counters)
.unwrap();
resolutions.insert(
Principal(1),
Resolution::Negative {
deadline: t(10),
next_refetch: t(10),
watermark: None,
},
);
assert!(!resolutions.publish(
Publications::Refreshed(vec![stale.read(|| PublishableSnapshotUpdate::Unknown {
principal: Principal(1),
until: t(10)
})]),
t(0),
&counters
));
assert_eq!(resolutions.unresolved(t(0)), 2);
assert!(resolutions.by_principal.is_empty());
assert!(resolutions.negative.is_empty());
assert!(resolutions.refetch.is_empty());
assert_eq!(counters.snapshot().publication_failures, 1);
assert_eq!(counters.snapshot().refresh_failures, 0);
assert!(map.get(&Principal(1)).is_none());
}
#[test]
fn control_wakeup_tracks_future_positive_deadline_once() {
let principal = Principal(1);
let mut resolutions = Resolutions::new([Principal(1)]);
resolutions.insert(
principal,
Resolution::Present {
deadline: t(15),
generation: Generation(1),
},
);
assert_eq!(
resolutions.next_control_wakeup(t(10)),
std::time::Duration::from_secs(5)
);
resolutions.insert(
principal,
Resolution::Present {
deadline: t(10),
generation: Generation(1),
},
);
assert_eq!(
resolutions.next_control_wakeup(t(10)),
std::time::Duration::from_secs(3_600)
);
}
#[test]
fn control_wakeup_tracks_negative_refetch_deadline() {
let principal = Principal(1);
let mut resolutions = Resolutions::new([Principal(1)]);
resolutions.insert(
principal,
Resolution::Negative {
deadline: t(15),
next_refetch: t(15),
watermark: None,
},
);
assert_eq!(
resolutions.next_control_wakeup(t(10)),
std::time::Duration::from_secs(5)
);
assert_eq!(
resolutions.next_control_wakeup(t(15)),
std::time::Duration::ZERO
);
}
mod naive {
use super::*;
pub(super) fn unresolved(
principals: &[Principal],
resolutions: &HashMap<Principal, Resolution>,
now: jiff::Timestamp,
) -> usize {
principals
.iter()
.filter(|principal| {
!resolutions
.get(principal)
.is_some_and(|resolution| now < resolution.deadline())
})
.count()
}
pub(super) fn next_readiness_check(
resolutions: &HashMap<Principal, Resolution>,
now: jiff::Timestamp,
) -> std::time::Duration {
#[allow(
clippy::disallowed_methods,
reason = "reduces deadlines to a minimum; a minimum does not depend on the order it is taken in"
)]
resolutions
.values()
.map(|resolution| resolution.deadline())
.filter(|deadline| *deadline > now)
.map(|deadline| deadline.duration_since(now).as_nanos())
.min()
.and_then(|nanos| u64::try_from(nanos).ok())
.map(std::time::Duration::from_nanos)
.unwrap_or_else(|| std::time::Duration::from_secs(3_600))
}
pub(super) fn next_control_wakeup(
resolutions: &HashMap<Principal, Resolution>,
now: jiff::Timestamp,
) -> std::time::Duration {
#[allow(
clippy::disallowed_methods,
reason = "reduces deadlines to a minimum; a minimum does not depend on the order it is taken in"
)]
resolutions
.values()
.filter_map(|resolution| match resolution {
Resolution::Present { deadline, .. } if *deadline > now => Some(*deadline),
Resolution::Present { .. } => None,
Resolution::Negative { next_refetch, .. } => Some(*next_refetch),
})
.map(|deadline| {
if deadline <= now {
std::time::Duration::ZERO
} else {
let nanos = deadline.duration_since(now).as_nanos();
u64::try_from(nanos)
.map(std::time::Duration::from_nanos)
.unwrap_or(std::time::Duration::MAX)
}
})
.min()
.unwrap_or_else(|| std::time::Duration::from_secs(3_600))
}
pub(super) fn due_for_refetch(
resolutions: &HashMap<Principal, Resolution>,
now: jiff::Timestamp,
limit: usize,
) -> Vec<Principal> {
#[allow(
clippy::disallowed_methods,
reason = "sorted below before the limit is applied, so which principals a bounded refetch takes is decided by deadline"
)]
let mut due: Vec<_> = resolutions
.iter()
.filter_map(|(principal, resolution)| match resolution {
Resolution::Negative {
next_refetch,
deadline,
..
} if *next_refetch <= now => Some((*next_refetch, *deadline, *principal)),
_ => None,
})
.collect();
due.sort_unstable_by_key(|(next_refetch, _, principal)| (*next_refetch, *principal));
due.into_iter()
.take(limit)
.map(|(_, _, principal)| principal)
.collect()
}
}
#[derive(Debug, Clone, Copy)]
enum Step {
Present { principal: u8, deadline: i64 },
Negative { principal: u8, deadline: i64 },
BackOff { principal: u8, retry_in: i64 },
Advance { seconds: i64 },
}
const TRACKED: usize = 6;
fn step() -> impl Strategy<Value = Step> {
prop_oneof![
(0..TRACKED as u8, 0i64..400).prop_map(|(principal, deadline)| Step::Present {
principal,
deadline
}),
(0..TRACKED as u8, 0i64..400).prop_map(|(principal, deadline)| Step::Negative {
principal,
deadline
}),
(0..TRACKED as u8, 0i64..200).prop_map(|(principal, retry_in)| Step::BackOff {
principal,
retry_in
}),
(0i64..50).prop_map(|seconds| Step::Advance { seconds }),
]
}
proptest! {
#[test]
fn indexed_resolutions_answer_exactly_what_scanning_answered(
steps in proptest::collection::vec(step(), 1..60),
) {
let principals: Vec<Principal> =
(0..TRACKED as u128).map(Principal).collect();
let mut indexed = Resolutions::new(principals.iter().copied());
let mut reference: HashMap<Principal, Resolution> = HashMap::new();
let mut now = t(0);
for step in steps {
match step {
Step::Present { principal, deadline } => {
let principal = Principal(u128::from(principal));
let resolution = Resolution::Present {
deadline: t(deadline),
generation: Generation(1),
};
indexed.insert(principal, resolution);
reference.insert(principal, resolution);
}
Step::Negative { principal, deadline } => {
let principal = Principal(u128::from(principal));
let resolution = Resolution::Negative {
deadline: t(deadline),
next_refetch: t(deadline),
watermark: None,
};
indexed.insert(principal, resolution);
reference.insert(principal, resolution);
}
Step::BackOff { principal, retry_in } => {
let principal = Principal(u128::from(principal));
let retry_at = t(now.as_second() + retry_in);
indexed.back_off(principal, retry_at);
if let Some(Resolution::Negative { next_refetch, .. }) =
reference.get_mut(&principal)
{
*next_refetch = retry_at;
}
}
Step::Advance { seconds } => {
now = t(now.as_second() + seconds);
}
}
prop_assert_eq!(
indexed.unresolved(now),
naive::unresolved(&principals, &reference, now),
"unresolved disagreed at {:?}", now
);
prop_assert_eq!(
indexed.next_readiness_check(now),
naive::next_readiness_check(&reference, now),
"next_readiness_check disagreed at {:?}", now
);
prop_assert_eq!(
indexed.next_control_wakeup(now),
naive::next_control_wakeup(&reference, now),
"next_control_wakeup disagreed at {:?}", now
);
for limit in [1usize, 3, usize::MAX] {
let mut due = indexed.due_for_refetch(now, limit);
let mut expected = naive::due_for_refetch(&reference, now, limit);
due.sort_unstable();
expected.sort_unstable();
prop_assert_eq!(
due,
expected,
"due_for_refetch disagreed at {:?} under limit {}", now, limit
);
}
}
}
}
#[test]
fn a_resolution_inserted_expired_is_never_live() {
let mut resolutions = Resolutions::new([Principal(0)]);
resolutions.insert(
Principal(0),
Resolution::Present {
deadline: t(5),
generation: Generation(1),
},
);
assert_eq!(resolutions.unresolved(t(10)), 1);
assert_eq!(resolutions.unresolved(t(10)), 1);
assert_eq!(resolutions.unresolved(t(20)), 1);
}
#[test]
fn re_resolving_a_principal_does_not_double_count_it() {
let mut resolutions = Resolutions::new([Principal(0), Principal(1)]);
for deadline in [t(50), t(60), t(60), t(70)] {
resolutions.insert(
Principal(0),
Resolution::Present {
deadline,
generation: Generation(1),
},
);
}
assert_eq!(
resolutions.unresolved(t(10)),
1,
"one principal resolved, one still outstanding"
);
assert_eq!(resolutions.next_readiness_check(t(10)), secs(60));
}
#[test]
fn a_positive_replacing_a_negative_clears_both_of_its_indexes() {
let mut resolutions = Resolutions::new([Principal(0)]);
resolutions.insert(
Principal(0),
Resolution::Negative {
deadline: t(30),
next_refetch: t(30),
watermark: None,
},
);
resolutions.insert(
Principal(0),
Resolution::Present {
deadline: t(90),
generation: Generation(2),
},
);
assert_eq!(resolutions.unresolved(t(40)), 0, "the positive is live");
assert!(
resolutions.due_for_refetch(t(40), usize::MAX).is_empty(),
"the superseded negative must not still ask to be refetched"
);
assert_eq!(resolutions.next_control_wakeup(t(40)), secs(50));
}
#[test]
fn a_due_population_is_refetched_in_bounded_waves() {
let principals: Vec<Principal> = (0..5).map(Principal).collect();
let mut resolutions = Resolutions::new(principals.iter().copied());
for (offset, principal) in principals.iter().enumerate() {
resolutions.insert(
*principal,
Resolution::Negative {
deadline: t(10 + offset as i64),
next_refetch: t(10 + offset as i64),
watermark: Some(Watermark::Revoked(Generation(9))),
},
);
}
assert_eq!(
resolutions.due_for_refetch(t(100), 2),
vec![Principal(0), Principal(1)],
"one wave, and the earliest deadlines lead it"
);
assert_eq!(
resolutions.due_for_refetch(t(100), usize::MAX).len(),
5,
"capping a wave must not retire the rest: they are still due"
);
}
#[test]
fn lag_recovery_covers_the_negatives_a_sweep_skips() {
let mut resolutions = Resolutions::new([Principal(0), Principal(1)]);
resolutions.insert(
Principal(0),
Resolution::Present {
deadline: t(90),
generation: Generation(2),
},
);
resolutions.insert(
Principal(1),
Resolution::Negative {
deadline: t(3_600),
next_refetch: t(3_600),
watermark: Some(Watermark::Revoked(Generation(9))),
},
);
assert_eq!(
resolutions.due_for_sweep(),
vec![Principal(0)],
"a routine sweep leaves the tombstone to its own schedule"
);
assert_eq!(
resolutions.all_tracked(),
vec![Principal(0), Principal(1)],
"lag recovery refetches everything, tombstones included"
);
}
#[test]
fn an_expired_resolution_keeps_its_generation() {
let mut resolutions = Resolutions::new([Principal(0)]);
resolutions.insert(
Principal(0),
Resolution::Negative {
deadline: t(30),
next_refetch: t(30),
watermark: Some(Watermark::Revoked(Generation(7))),
},
);
assert_eq!(resolutions.unresolved(t(100)), 1, "expired");
assert_eq!(
resolutions.watermark_of(Principal(0)),
Some(Watermark::Revoked(Generation(7))),
"the watermark must survive the expiry, provenance included"
);
}
fn secs(seconds: u64) -> std::time::Duration {
std::time::Duration::from_secs(seconds)
}
#[test]
fn deadline_helpers_are_exact_in_the_supported_domain() {
let config = SnapshotManagerConfig {
principals: TrackedPrincipals::Fixed(vec![Principal(1)]),
refresh_interval: std::time::Duration::from_secs(60),
unknown_ttl: SignedDuration::from_secs(30),
revoked_ttl: SignedDuration::from_secs(3_600),
retry_backoff: std::time::Duration::from_secs(2),
max_concurrent_fetches: 1,
fetch_timeout: std::time::Duration::from_secs(5),
enumeration_timeout: std::time::Duration::from_secs(30),
};
assert_eq!(
negative_deadline(t(10), &config, NegativeKind::Unknown),
t(40)
);
assert_eq!(
negative_deadline(t(10), &config, NegativeKind::Revoked),
t(3_610)
);
assert_eq!(after_std(t(10), config.retry_backoff), t(12));
}
}