Skip to main content

arete_server/snapshot/
mod.rs

1//! Opt-in periodic snapshots of in-memory server state (VM entity tables +
2//! projection caches) and a restore path that rehydrates on startup, so a
3//! restarted server comes back with its history instead of starting empty.
4//!
5//! arete-server owns all snapshot logic. The generated runtime's only
6//! responsibilities are (1) registering the `VmContext`/`SlotTracker` it
7//! creates via [`register_runtime`] and (2) hydrating from a restored blob via
8//! [`take_restored`] before connecting to Yellowstone. Those hooks resolve
9//! through a task-local [`SnapshotRuntime`], so multiple servers embedded in
10//! one process cannot consume or replace each other's snapshot state. A stack
11//! built with an older arete-macros simply never registers a VM; snapshots
12//! stay disabled with a warning.
13//!
14//! Consistency cut: every generated VM update holds a shared snapshot barrier
15//! guard until its mutation batch has been applied by the projector. Snapshot
16//! capture takes the exclusive guard before dumping either side, so the VM,
17//! projection caches, and resume watermark all describe the same processing
18//! cut. On restore the stream replays from that watermark; the snapshotted
19//! version trackers drop the overlap.
20
21pub mod envelope;
22#[cfg(feature = "snapshot-object-store")]
23pub mod object;
24pub mod store;
25
26pub use envelope::{SnapshotContract, SnapshotHeader, SnapshotPayload};
27#[cfg(feature = "snapshot-object-store")]
28pub use object::ObjectSnapshotStore;
29pub use store::{FsStore, SnapshotStore};
30
31use crate::cache::EntityCache;
32use crate::health::SlotTracker;
33use crate::mutation_batch::MutationBatch;
34use crate::view::ViewIndex;
35use anyhow::{Context, Result};
36use arete_interpreter::snapshot::{VmSnapshot, SNAPSHOT_FORMAT_VERSION};
37use arete_interpreter::vm::VmContext;
38use serde::Serialize;
39use sha2::{Digest, Sha256};
40use std::collections::{BTreeMap, BTreeSet, HashMap};
41use std::future::Future;
42use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
43use std::sync::{Arc, Mutex as StdMutex};
44use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
45use tokio::sync::{mpsc, OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock};
46use tracing::{debug, info, info_span, warn, Instrument};
47
48/// Rough Solana slot duration, used only to convert snapshot age into an
49/// estimated slot distance for the staleness clamp.
50const ESTIMATED_SLOT_MILLIS: u64 = 200;
51/// How long a snapshot cycle waits for in-flight VM updates and their queued
52/// projection batches to finish.
53const CONSISTENCY_CUT_TIMEOUT: Duration = Duration::from_secs(10);
54const STATE_CONTRACT_SCHEMA_V1: &str = "arete.snapshot-state-contract/v1";
55const PROJECTION_CONTRACT_SCHEMA_V1: &str = "arete.snapshot-projection-contract/v1";
56
57/// Configuration for state snapshots. Disabled by default; enable via
58/// `ServerBuilder::snapshots(...)` or `ARETE_SNAPSHOT_*` env vars.
59#[derive(Clone, Debug)]
60pub struct SnapshotConfig {
61    /// Master opt-in.
62    pub enabled: bool,
63    /// Where blobs live: `file:///var/lib/arete/snapshots`, a plain path, or
64    /// (with the `snapshot-object-store` feature) `s3://`/`gs://`/`az://`.
65    pub url: Option<String>,
66    /// Periodic snapshot cadence.
67    pub interval: Duration,
68    /// Retained snapshots; older ones are pruned after each write.
69    pub keep: usize,
70    /// Take a final snapshot on SIGTERM/SIGINT before exit.
71    pub snapshot_on_shutdown: bool,
72    /// Skip a periodic cycle when fewer batches were applied since the last
73    /// snapshot (quiet stacks snapshot rarely).
74    pub min_mutations: u64,
75    /// If the watermark's estimated lag is greater than this many slots,
76    /// including lag at snapshot time plus file age, hydrate state but start
77    /// the stream live instead of resuming from the watermark.
78    pub max_resume_age_slots: u64,
79    /// `/ready` stays 503 after a watermark resume until the projector is
80    /// within this many slots of the observed tip...
81    pub ready_max_lag_slots: u64,
82    /// ...or until this much time has passed (guards quiet stacks, where the
83    /// watermark never advances because nothing happens on-chain).
84    pub ready_max_hold: Duration,
85    /// Bytecode hashes from pre-contract snapshots that an operator has
86    /// explicitly approved for one state-only migration. These snapshots are
87    /// hydrated after structural state-id remapping and always start live.
88    pub legacy_bytecode_hashes: BTreeSet<String>,
89}
90
91impl Default for SnapshotConfig {
92    fn default() -> Self {
93        Self {
94            enabled: false,
95            url: None,
96            interval: Duration::from_secs(60),
97            keep: 4,
98            snapshot_on_shutdown: true,
99            min_mutations: 1,
100            // ~10 minutes of slots: conservative vs. typical provider
101            // `from_slot` replay windows (in-cluster richat rings are far
102            // more generous; raw Triton is minutes).
103            max_resume_age_slots: 1_500,
104            ready_max_lag_slots: 50,
105            ready_max_hold: Duration::from_secs(60),
106            legacy_bytecode_hashes: BTreeSet::new(),
107        }
108    }
109}
110
111impl SnapshotConfig {
112    /// Load snapshot settings from `ARETE_SNAPSHOT_*` env vars. Snapshots stay
113    /// disabled unless `ARETE_SNAPSHOT_ENABLED=true`.
114    pub fn from_env() -> Result<Self> {
115        let mut config = Self::default();
116        config.enabled = crate::config::env_bool("ARETE_SNAPSHOT_ENABLED")?.unwrap_or(false);
117        config.url = std::env::var("ARETE_SNAPSHOT_URL")
118            .ok()
119            .filter(|value| !value.trim().is_empty());
120        config.interval = Duration::from_secs(
121            crate::config::env_parse("ARETE_SNAPSHOT_INTERVAL_SECS")?
122                .unwrap_or(config.interval.as_secs()),
123        );
124        config.keep = crate::config::env_parse("ARETE_SNAPSHOT_KEEP")?.unwrap_or(config.keep);
125        config.snapshot_on_shutdown = crate::config::env_bool("ARETE_SNAPSHOT_ON_SHUTDOWN")?
126            .unwrap_or(config.snapshot_on_shutdown);
127        config.min_mutations = crate::config::env_parse("ARETE_SNAPSHOT_MIN_MUTATIONS")?
128            .unwrap_or(config.min_mutations);
129        config.max_resume_age_slots =
130            crate::config::env_parse("ARETE_SNAPSHOT_MAX_RESUME_AGE_SLOTS")?
131                .unwrap_or(config.max_resume_age_slots);
132        config.ready_max_lag_slots =
133            crate::config::env_parse("ARETE_SNAPSHOT_READY_MAX_LAG_SLOTS")?
134                .unwrap_or(config.ready_max_lag_slots);
135        config.ready_max_hold = Duration::from_secs(
136            crate::config::env_parse("ARETE_SNAPSHOT_READY_MAX_HOLD_SECS")?
137                .unwrap_or(config.ready_max_hold.as_secs()),
138        );
139        config.legacy_bytecode_hashes = std::env::var("ARETE_SNAPSHOT_LEGACY_BYTECODE_HASHES")
140            .ok()
141            .into_iter()
142            .flat_map(|value| {
143                value
144                    .split(',')
145                    .map(str::trim)
146                    .filter(|hash| !hash.is_empty())
147                    .map(str::to_ascii_lowercase)
148                    .collect::<Vec<_>>()
149            })
150            .collect();
151        config.validate()?;
152        Ok(config)
153    }
154
155    pub fn validate(&self) -> Result<()> {
156        if self.enabled && self.url.as_deref().is_none_or(|url| url.trim().is_empty()) {
157            anyhow::bail!("snapshots are enabled but ARETE_SNAPSHOT_URL is not set");
158        }
159        if self.enabled && (self.interval.is_zero() || self.keep == 0) {
160            anyhow::bail!("snapshot interval and keep count must be greater than zero");
161        }
162        for hash in &self.legacy_bytecode_hashes {
163            if hash.len() != 64 || !hash.bytes().all(|byte| byte.is_ascii_hexdigit()) {
164                anyhow::bail!(
165                    "ARETE_SNAPSHOT_LEGACY_BYTECODE_HASHES contains invalid SHA-256 hash '{hash}'"
166                );
167            }
168        }
169        Ok(())
170    }
171}
172
173fn contract_hash<T: Serialize>(schema: &str, value: &T) -> SnapshotContract {
174    let canonical = arete_hash::canonicalize_jcs(value)
175        .expect("snapshot contracts contain only canonical JSON values");
176    let hash = hex::encode(Sha256::digest(canonical));
177    SnapshotContract {
178        schema: schema.to_string(),
179        hash,
180    }
181}
182
183fn state_contract(spec: &crate::Spec) -> Option<SnapshotContract> {
184    if spec.entity_specs.is_empty() {
185        return None;
186    }
187    let mut entities = Vec::with_capacity(spec.entity_specs.len());
188    for entity in &spec.entity_specs {
189        let mut indexes = entity
190            .identity
191            .lookup_indexes
192            .iter()
193            .map(|index| (&index.field_name, &index.temporal_field))
194            .collect::<Vec<_>>();
195        indexes.sort();
196        indexes.dedup();
197        let fields = entity
198            .field_mappings
199            .iter()
200            .map(|(path, field)| {
201                (
202                    path,
203                    serde_json::json!({
204                        "baseType": field.base_type,
205                        "integerKind": field.integer_kind,
206                        "isOptional": field.is_optional,
207                        "isArray": field.is_array,
208                        "innerType": field.inner_type,
209                        "resolvedType": field.resolved_type,
210                        "emit": field.emit,
211                    }),
212                )
213            })
214            .collect::<BTreeMap<_, _>>();
215        entities.push(serde_json::json!({
216            "name": entity.state_name,
217            "primaryKeys": entity.identity.primary_keys,
218            "lookupIndexes": indexes,
219            "fields": fields,
220        }));
221    }
222    entities.sort_by_key(|entity| entity["name"].as_str().unwrap_or_default().to_string());
223    Some(contract_hash(
224        STATE_CONTRACT_SCHEMA_V1,
225        &serde_json::json!({"entities": entities}),
226    ))
227}
228
229fn projection_contract(view_index: &ViewIndex) -> SnapshotContract {
230    contract_hash(PROJECTION_CONTRACT_SCHEMA_V1, &view_index.snapshot_specs())
231}
232
233fn state_ids_by_entity(spec: &crate::Spec) -> HashMap<String, u32> {
234    spec.bytecode
235        .entities
236        .iter()
237        .map(|(name, entity)| (name.clone(), entity.state_id))
238        .collect()
239}
240
241fn remap_snapshot_states(vm: &mut VmSnapshot, state_ids: &HashMap<String, u32>) -> Result<()> {
242    let states = std::mem::take(&mut vm.states);
243    for (_, table) in states {
244        let state_id = state_ids
245            .get(&table.entity_name)
246            .with_context(|| format!("snapshot contains unknown entity '{}'", table.entity_name))?;
247        if vm.states.insert(*state_id, table).is_some() {
248            anyhow::bail!("snapshot contains duplicate entity state for id {state_id}");
249        }
250    }
251    Ok(())
252}
253
254/// VM state handed from the restore path to the generated runtime, consumed
255/// exactly once via [`take_restored`].
256pub struct RestoredState {
257    pub vm: VmSnapshot,
258    /// `Some(slot)` to resume the Yellowstone stream from that slot; `None`
259    /// when the snapshot was too stale (state still hydrates, stream starts
260    /// live and account-derived state self-heals).
261    pub resume_watermark: Option<u64>,
262}
263
264#[derive(Clone)]
265struct RuntimeRegistration {
266    vm: Arc<StdMutex<VmContext>>,
267    slot_tracker: SlotTracker,
268}
269
270struct ResumeGate {
271    started: Instant,
272    max_lag_slots: u64,
273    max_hold: Duration,
274}
275
276/// Per-runtime barrier that keeps snapshot capture from splitting a VM update
277/// from the projection batch it produced.
278///
279/// Generated mutation producers enter the barrier in shared mode before
280/// touching the VM and transfer the guard to their [`MutationBatch`]. The
281/// projector releases it only after applying that batch. Snapshot capture
282/// enters in exclusive mode, which therefore waits for both in-flight parser
283/// work and queued projection work to finish.
284#[derive(Clone, Default)]
285pub struct SnapshotBarrier {
286    inner: Arc<RwLock<()>>,
287}
288
289impl SnapshotBarrier {
290    pub async fn enter_processing(&self) -> SnapshotProcessingGuard {
291        SnapshotProcessingGuard(self.inner.clone().read_owned().await)
292    }
293
294    async fn enter_snapshot(&self) -> OwnedRwLockWriteGuard<()> {
295        self.inner.clone().write_owned().await
296    }
297}
298
299/// Shared processing guard carried by a mutation batch until projection is
300/// complete. The inner guard is intentionally opaque outside arete-server.
301pub struct SnapshotProcessingGuard(#[allow(dead_code)] OwnedRwLockReadGuard<()>);
302
303impl std::fmt::Debug for SnapshotProcessingGuard {
304    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
305        formatter.write_str("SnapshotProcessingGuard")
306    }
307}
308
309#[derive(Default)]
310struct SnapshotRuntimeState {
311    registered: StdMutex<Option<RuntimeRegistration>>,
312    restored: StdMutex<Option<RestoredState>>,
313    resume_gate: StdMutex<Option<ResumeGate>>,
314    processing_barrier: SnapshotBarrier,
315    /// Highest slot among mutation batches this runtime's projector has
316    /// applied. This is the safe `from_slot` resume point (`SlotTracker` is
317    /// not: it follows the raw slot subscription, not parser progress).
318    resume_watermark: AtomicU64,
319    applied_batches: AtomicU64,
320}
321
322/// Per-server snapshot coordination shared by its parser, projector, snapshot
323/// manager, and readiness endpoint.
324///
325/// The generated parser hooks use [`scope`](Self::scope) so their existing
326/// argument-free calls cannot accidentally bind to another server running in
327/// the same process.
328#[derive(Clone, Default)]
329pub struct SnapshotRuntime {
330    state: Arc<SnapshotRuntimeState>,
331}
332
333tokio::task_local! {
334    static ACTIVE_SNAPSHOT_RUNTIME: SnapshotRuntime;
335}
336
337impl SnapshotRuntime {
338    /// Run a generated parser future with this server's snapshot state.
339    pub async fn scope<F>(&self, future: F) -> F::Output
340    where
341        F: Future,
342    {
343        ACTIVE_SNAPSHOT_RUNTIME.scope(self.clone(), future).await
344    }
345
346    /// Associate the parser's VM and slot tracker with this server only.
347    pub fn register_runtime(
348        &self,
349        vm: Arc<StdMutex<VmContext>>,
350        slot_tracker: SlotTracker,
351    ) -> SnapshotBarrier {
352        let mut registered = self.state.registered.lock().unwrap();
353        if registered.is_some() {
354            debug!("Snapshot runtime registration replaced");
355        }
356        *registered = Some(RuntimeRegistration { vm, slot_tracker });
357        self.state.processing_barrier.clone()
358    }
359
360    /// Consume this server's restored VM state exactly once.
361    pub fn take_restored(&self) -> Option<RestoredState> {
362        self.state.restored.lock().unwrap().take()
363    }
364
365    /// Record a batch applied by this server's projector.
366    pub(crate) fn record_applied_batch(&self, slot: Option<u64>) {
367        self.state.applied_batches.fetch_add(1, Ordering::Relaxed);
368        if let Some(slot) = slot {
369            self.state
370                .resume_watermark
371                .fetch_max(slot, Ordering::Relaxed);
372        }
373    }
374
375    /// Returns `true` unless this server's watermark resume is still catching
376    /// up to its observed slot tip.
377    pub fn resume_gate_ready(&self) -> bool {
378        let mut gate_slot = self.state.resume_gate.lock().unwrap();
379        let Some(gate) = gate_slot.as_ref() else {
380            return true;
381        };
382        if gate.started.elapsed() >= gate.max_hold {
383            info!("Snapshot resume readiness gate released (max hold reached)");
384            *gate_slot = None;
385            return true;
386        }
387        let tip = self
388            .state
389            .registered
390            .lock()
391            .unwrap()
392            .as_ref()
393            .map(|registration| registration.slot_tracker.get())
394            .unwrap_or(0);
395        let applied = self.state.resume_watermark.load(Ordering::Relaxed);
396        if tip > 0 && tip.saturating_sub(applied) <= gate.max_lag_slots {
397            info!(tip, applied, "Snapshot resume caught up; marking ready");
398            *gate_slot = None;
399            return true;
400        }
401        false
402    }
403}
404
405/// Called by the generated runtime after it creates its `VmContext` and
406/// `SlotTracker`, so the snapshot manager can dump them later.
407pub fn register_runtime(
408    vm: Arc<StdMutex<VmContext>>,
409    slot_tracker: SlotTracker,
410) -> Option<SnapshotBarrier> {
411    match ACTIVE_SNAPSHOT_RUNTIME.try_with(|runtime| runtime.register_runtime(vm, slot_tracker)) {
412        Ok(barrier) => Some(barrier),
413        Err(_) => {
414            debug!("Snapshot runtime registration ignored (snapshots disabled)");
415            None
416        }
417    }
418}
419
420/// Called by the generated runtime before connecting: returns the restored VM
421/// state (if any) exactly once.
422pub fn take_restored() -> Option<RestoredState> {
423    ACTIVE_SNAPSHOT_RUNTIME
424        .try_with(SnapshotRuntime::take_restored)
425        .ok()
426        .flatten()
427}
428
429/// Where the generated Yellowstone runtime should resume its stream.
430///
431/// `Option<u64>` cannot express this: "no checkpoint to resume from" and
432/// "gave up on the checkpoint" are both `None`, and only the second one
433/// loses data. Naming them apart is what lets the caller mark the hole.
434#[derive(Debug, Clone, Copy, PartialEq, Eq)]
435pub enum ReconnectPosition {
436    /// Resume from this slot; nothing is lost.
437    Slot(u64),
438    /// Start live because nothing has been processed yet.
439    Live,
440    /// Start live after abandoning a checkpoint the provider would not serve.
441    /// Every slot between `abandoned` and the live tip is lost.
442    LiveAfterGap { abandoned: u64 },
443}
444
445impl ReconnectPosition {
446    /// The `from_slot` to put on the subscription request.
447    pub fn from_slot(self) -> Option<u64> {
448        match self {
449            Self::Slot(slot) => Some(slot),
450            Self::Live | Self::LiveAfterGap { .. } => None,
451        }
452    }
453}
454
455/// Select a reconnect checkpoint for the generated Yellowstone runtime.
456///
457/// A restored replay never falls back to live: retries advance only to slots
458/// the main parser stream has finished processing. Without a restored replay,
459/// repeated short-lived connections eventually give up on the checkpoint —
460/// unless `live_fallback_attempts` is `None`, which refuses to trade data for
461/// availability.
462#[doc(hidden)]
463pub fn select_reconnect_from_slot(
464    restored_watermark: Option<u64>,
465    processed_watermark: u64,
466    attempt: u32,
467    live_fallback_attempts: Option<u32>,
468) -> ReconnectPosition {
469    if let Some(restored_watermark) = restored_watermark {
470        return ReconnectPosition::Slot(restored_watermark.max(processed_watermark));
471    }
472    if processed_watermark == 0 {
473        // Nothing has been processed, so starting live loses nothing.
474        return ReconnectPosition::Live;
475    }
476    match live_fallback_attempts {
477        Some(limit) if attempt >= limit => ReconnectPosition::LiveAfterGap {
478            abandoned: processed_watermark,
479        },
480        _ => ReconnectPosition::Slot(processed_watermark),
481    }
482}
483
484fn now_epoch_ms() -> u64 {
485    SystemTime::now()
486        .duration_since(UNIX_EPOCH)
487        .unwrap()
488        .as_millis() as u64
489}
490
491/// What kicked off a snapshot cycle.
492#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
493#[serde(rename_all = "snake_case")]
494pub enum SnapshotTrigger {
495    Periodic,
496    Shutdown,
497}
498
499/// Owns the store plus everything needed to dump and restore state. Created by
500/// `Runtime::run` when snapshots are enabled.
501pub struct SnapshotService {
502    config: SnapshotConfig,
503    store: Arc<dyn SnapshotStore>,
504    runtime: SnapshotRuntime,
505    bytecode_hash: String,
506    state_contract: Option<SnapshotContract>,
507    projection_contract: SnapshotContract,
508    state_ids: HashMap<String, u32>,
509    program_ids: Vec<String>,
510    entity_cache: EntityCache,
511    journal: Arc<crate::journal::EventJournal>,
512    batches_at_last_snapshot: AtomicU64,
513    warned_missing_vm: AtomicBool,
514}
515
516impl SnapshotService {
517    /// Build the store, then attempt a restore (any failure logs a warning
518    /// and cold-starts — restore problems must never block startup).
519    pub async fn initialize(
520        config: SnapshotConfig,
521        spec: &crate::Spec,
522        entity_cache: EntityCache,
523        view_index: &ViewIndex,
524        journal: Arc<crate::journal::EventJournal>,
525        _mutations_tx: mpsc::Sender<MutationBatch>,
526    ) -> Result<Arc<Self>> {
527        let url = config
528            .url
529            .clone()
530            .context("snapshots are enabled but no snapshot URL is configured")?;
531        let store = store::store_from_url(&url)?;
532
533        let mut program_ids = spec.program_ids.clone();
534        program_ids.sort();
535
536        let service = Arc::new(Self {
537            config,
538            store,
539            runtime: SnapshotRuntime::default(),
540            bytecode_hash: spec.bytecode.fingerprint(),
541            state_contract: state_contract(spec),
542            projection_contract: projection_contract(view_index),
543            state_ids: state_ids_by_entity(spec),
544            program_ids,
545            entity_cache,
546            journal,
547            batches_at_last_snapshot: AtomicU64::new(0),
548            warned_missing_vm: AtomicBool::new(false),
549        });
550        info!(
551            store = %service.store.describe(),
552            interval_secs = service.config.interval.as_secs(),
553            keep = service.config.keep,
554            "State snapshots enabled"
555        );
556
557        match service.restore(view_index).await {
558            Ok(true) => {}
559            Ok(false) => info!("No usable snapshot found; starting cold"),
560            Err(err) => warn!(
561                error = format!("{err:#}"),
562                "Failed to restore snapshot; starting cold"
563            ),
564        }
565        Ok(service)
566    }
567
568    pub fn config(&self) -> &SnapshotConfig {
569        &self.config
570    }
571
572    /// Return the per-server coordination handle that must be shared with the
573    /// matching parser, projector, and readiness endpoint.
574    pub fn runtime(&self) -> SnapshotRuntime {
575        self.runtime.clone()
576    }
577
578    /// Load and validate the latest snapshot, hydrate the projection caches,
579    /// and stash the VM portion for the generated runtime. Returns whether a
580    /// snapshot was applied.
581    async fn restore(&self, view_index: &ViewIndex) -> Result<bool> {
582        let Some((name, bytes)) = self.store.load_latest().await? else {
583            return Ok(false);
584        };
585
586        let header = envelope::decode_header(&bytes)
587            .with_context(|| format!("snapshot {name} has an unreadable header"))?;
588
589        if header.format_version != SNAPSHOT_FORMAT_VERSION {
590            warn!(
591                snapshot = %name,
592                found = header.format_version,
593                expected = SNAPSHOT_FORMAT_VERSION,
594                "Snapshot format version mismatch; discarding (cold start)"
595            );
596            return Ok(false);
597        }
598        let exact_bytecode = header.bytecode_hash == self.bytecode_hash;
599        let matching_contracts = self.state_contract.is_some()
600            && header.state_contract == self.state_contract
601            && header.projection_contract.as_ref() == Some(&self.projection_contract);
602        let approved_legacy = !exact_bytecode
603            && header.state_contract.is_none()
604            && header.projection_contract.is_none()
605            && self
606                .config
607                .legacy_bytecode_hashes
608                .contains(&header.bytecode_hash);
609        let legacy_migration = approved_legacy;
610        if !exact_bytecode && !matching_contracts && !approved_legacy {
611            warn!(
612                snapshot = %name,
613                "Snapshot was taken by an incompatible stack build; discarding (cold start)"
614            );
615            return Ok(false);
616        }
617        let mut snapshot_program_ids = header.program_ids.clone();
618        snapshot_program_ids.sort();
619        if snapshot_program_ids != self.program_ids {
620            warn!(
621                snapshot = %name,
622                "Snapshot program ids do not match this server; discarding (cold start)"
623            );
624            return Ok(false);
625        }
626
627        let mut payload = tokio::task::spawn_blocking(move || envelope::decode_payload(&bytes))
628            .await
629            .context("snapshot decode task panicked")?
630            .with_context(|| format!("snapshot {name} has an unreadable payload"))?;
631        if !exact_bytecode {
632            remap_snapshot_states(&mut payload.vm, &self.state_ids)
633                .with_context(|| format!("snapshot {name} state contract is incompatible"))?;
634        }
635        if legacy_migration {
636            // A pre-contract snapshot cannot prove that its materialized views
637            // still match the current projections. Preserve only durable VM
638            // state and let live input rebuild every projection cache.
639            payload.entity_cache.clear();
640            // Retained frames are published view output, shaped by the same
641            // projections, so the same doubt applies — and replaying stale
642            // frames is worse than a stale cache, because consumers keep them.
643            payload.journal = Default::default();
644        }
645
646        let cached_views = payload.entity_cache.len();
647        let cached_entities: usize = payload
648            .entity_cache
649            .iter()
650            .map(|(_, entries)| entries.len())
651            .sum();
652        let retained_events: usize = payload
653            .journal
654            .views
655            .values()
656            .map(|view| view.records.len())
657            .sum();
658        self.entity_cache.hydrate(payload.entity_cache).await;
659        // Only a shutdown snapshot is exact; see `EventJournal::hydrate`.
660        let exact_offsets = header.trigger == Some(SnapshotTrigger::Shutdown);
661        self.journal.hydrate(payload.journal, exact_offsets).await;
662        rebuild_sorted_caches(view_index, &self.entity_cache).await;
663
664        // Even when the stream starts live, the watermark seeds the applied
665        // position: the hydrated state already contains everything up to it.
666        self.runtime
667            .state
668            .resume_watermark
669            .fetch_max(header.resume_watermark, Ordering::Relaxed);
670
671        let age_ms = now_epoch_ms().saturating_sub(header.created_at_epoch_ms);
672        let estimated_age_slots = age_ms / ESTIMATED_SLOT_MILLIS;
673        // File age alone is insufficient: a shutdown checkpoint is freshly
674        // written even when a quiet program's last applied update is already
675        // far behind the observed chain tip. Account for both the lag already
676        // present at the consistency cut and the time elapsed since it.
677        let watermark_lag_at_snapshot =
678            header.observed_slot.saturating_sub(header.resume_watermark);
679        let estimated_resume_lag = watermark_lag_at_snapshot.saturating_add(estimated_age_slots);
680        let resume_watermark = if exact_bytecode
681            && header.resume_watermark > 0
682            && estimated_resume_lag <= self.config.max_resume_age_slots
683        {
684            Some(header.resume_watermark)
685        } else {
686            if header.resume_watermark > 0 && exact_bytecode {
687                warn!(
688                    resume_watermark = header.resume_watermark,
689                    observed_slot = header.observed_slot,
690                    watermark_lag_at_snapshot,
691                    estimated_age_slots,
692                    estimated_resume_lag,
693                    max_resume_age_slots = self.config.max_resume_age_slots,
694                    "Snapshot watermark is outside the resume window; hydrating state but \
695                     starting the stream live. Account-derived state self-heals from \
696                     full account writes; only instruction events in the gap are missed."
697                );
698            }
699            if !exact_bytecode {
700                warn!(
701                    snapshot = %name,
702                    matching_contracts,
703                    approved_legacy,
704                    legacy_state_only = legacy_migration,
705                    "Hydrated snapshot state from different bytecode; starting live"
706                );
707            }
708            None
709        };
710
711        if resume_watermark.is_none() {
712            // The stream starts live, so events between the retained tape and
713            // the first live append are lost. Offsets stay dense across that
714            // hole, which would present it to a consumer as an unbroken
715            // continuation — mark it so a replay across it is refused instead.
716            self.journal.mark_gap().await;
717        }
718
719        if resume_watermark.is_some() {
720            *self.runtime.state.resume_gate.lock().unwrap() = Some(ResumeGate {
721                started: Instant::now(),
722                max_lag_slots: self.config.ready_max_lag_slots,
723                max_hold: self.config.ready_max_hold,
724            });
725        }
726
727        info!(
728            snapshot = %name,
729            vm_entities = payload.vm.total_entries(),
730            cached_views,
731            cached_entities,
732            retained_events,
733            resume_watermark = header.resume_watermark,
734            resuming = resume_watermark.is_some(),
735            age_secs = age_ms / 1_000,
736            "Restored state from snapshot"
737        );
738
739        *self.runtime.state.restored.lock().unwrap() = Some(RestoredState {
740            vm: payload.vm,
741            resume_watermark,
742        });
743        Ok(true)
744    }
745
746    /// Spawn the periodic snapshot task.
747    pub fn spawn(self: &Arc<Self>) -> tokio::task::JoinHandle<()> {
748        let service = Arc::clone(self);
749        tokio::spawn(
750            async move {
751                let mut interval = tokio::time::interval(service.config.interval);
752                interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
753                // The first tick fires immediately; skip it so the first
754                // snapshot lands one full interval after startup.
755                interval.tick().await;
756                loop {
757                    interval.tick().await;
758                    if let Err(err) = service.snapshot_now(SnapshotTrigger::Periodic).await {
759                        // Snapshotting must never take down a healthy server.
760                        warn!(
761                            error = format!("{err:#}"),
762                            "Snapshot cycle failed; will retry next interval"
763                        );
764                    }
765                }
766            }
767            .instrument(info_span!("snapshot.manager")),
768        )
769    }
770
771    /// Run one snapshot cycle. Returns `Ok(false)` when skipped (no VM
772    /// registered yet, or too few mutations since the last snapshot).
773    pub async fn snapshot_now(&self, trigger: SnapshotTrigger) -> Result<bool> {
774        let Some(registration) = self.runtime.state.registered.lock().unwrap().clone() else {
775            if !self.warned_missing_vm.swap(true, Ordering::Relaxed) {
776                warn!(
777                    "Snapshots are enabled but no VM has been registered; the stack \
778                     may have been built with an older arete-macros version"
779                );
780            }
781            return Ok(false);
782        };
783
784        let applied_batches = self.runtime.state.applied_batches.load(Ordering::Relaxed);
785        if trigger == SnapshotTrigger::Periodic {
786            let since_last = applied_batches
787                .saturating_sub(self.batches_at_last_snapshot.load(Ordering::Relaxed));
788            if since_last < self.config.min_mutations {
789                debug!(since_last, "Skipping snapshot cycle (too few mutations)");
790                return Ok(false);
791            }
792        }
793
794        // Wait for every in-flight VM update and its queued projection batch
795        // to finish, then block new updates until both sides have been dumped.
796        // Processing guards move with their mutation batches and are released
797        // by the projector only after cache application, so this exclusive
798        // guard establishes one exact cut without an enqueue race.
799        let consistency_guard = tokio::time::timeout(
800            CONSISTENCY_CUT_TIMEOUT,
801            self.runtime.state.processing_barrier.enter_snapshot(),
802        )
803        .await
804        .context("timed out waiting for a consistent VM/projection snapshot cut")?;
805
806        let dump_started = Instant::now();
807        let (vm_snapshot, resume_watermark) = {
808            let vm = registration
809                .vm
810                .lock()
811                .map_err(|_| anyhow::anyhow!("VM mutex poisoned"))?;
812            let resume_watermark = self.runtime.state.resume_watermark.load(Ordering::Relaxed);
813            let vm_snapshot = vm.dump();
814            (vm_snapshot, resume_watermark)
815        };
816        let vm_lock_held = dump_started.elapsed();
817        let observed_slot = registration.slot_tracker.get();
818        let entity_cache_dump = self.entity_cache.dump().await;
819        // Dumped inside the same consistency guard as the cache, so a restore
820        // can never leave the cache ahead of the tape.
821        let journal_dump = self.journal.dump().await;
822        if trigger == SnapshotTrigger::Shutdown {
823            // Publishing continues after the guard releases — the parser is
824            // aborted only once this snapshot is encoded and stored — so
825            // without this the file would not hold every offset that reached a
826            // subscriber, and the restore below would adopt its epoch anyway.
827            self.journal.seal();
828        }
829        let applied_batches = self.runtime.state.applied_batches.load(Ordering::Relaxed);
830        drop(consistency_guard);
831
832        let created_at_epoch_ms = now_epoch_ms();
833        let header = SnapshotHeader {
834            format_version: SNAPSHOT_FORMAT_VERSION,
835            bytecode_hash: self.bytecode_hash.clone(),
836            state_contract: self.state_contract.clone(),
837            projection_contract: Some(self.projection_contract.clone()),
838            program_ids: self.program_ids.clone(),
839            resume_watermark,
840            observed_slot,
841            created_at_epoch_ms,
842            trigger: Some(trigger),
843            entry_counts: vm_snapshot
844                .entry_counts()
845                .into_iter()
846                .chain(
847                    // Retained record counts are otherwise invisible after the
848                    // restore log line.
849                    // Derived from the dump rather than a second trip
850                    // through the journal's lock inside the cut.
851                    journal_dump.views.iter().map(|(view_id, view)| {
852                        (format!("journal:{view_id}"), view.records.len() as u64)
853                    }),
854                )
855                .collect(),
856        };
857        let payload = SnapshotPayload {
858            vm: vm_snapshot,
859            entity_cache: entity_cache_dump,
860            journal: journal_dump,
861        };
862        let bytes = tokio::task::spawn_blocking(move || envelope::encode(&header, &payload))
863            .await
864            .context("snapshot encode task panicked")??;
865
866        let name = store::snapshot_name(created_at_epoch_ms, resume_watermark);
867        self.store.write(&name, &bytes).await?;
868        if let Err(err) = self.store.prune(self.config.keep).await {
869            warn!(error = format!("{err:#}"), "Failed to prune old snapshots");
870        }
871        self.batches_at_last_snapshot
872            .store(applied_batches, Ordering::Relaxed);
873
874        info!(
875            snapshot = %name,
876            bytes = bytes.len(),
877            resume_watermark,
878            observed_slot,
879            vm_lock_ms = vm_lock_held.as_millis() as u64,
880            trigger = ?trigger,
881            "Snapshot written"
882        );
883        Ok(true)
884    }
885}
886
887/// Rebuild each derived `SortedViewCache` from the hydrated `EntityCache`,
888/// mirroring the projector's upsert path. Sorted caches are not persisted:
889/// they are derived state and MB-scale rebuilds are sub-millisecond.
890async fn rebuild_sorted_caches(view_index: &ViewIndex, entity_cache: &EntityCache) {
891    let sorted_caches = view_index.sorted_caches();
892    let max_entries = entity_cache.max_entities_per_view();
893    for spec in view_index.get_derived_views() {
894        let Some(source_view) = spec.source_view.as_ref() else {
895            continue;
896        };
897        let entities = entity_cache.get_all(source_view).await;
898        if entities.is_empty() {
899            continue;
900        }
901        let filter = spec
902            .pipeline
903            .as_ref()
904            .and_then(|pipeline| pipeline.filter.as_ref());
905        let mut caches = sorted_caches.write().await;
906        if let Some(cache) = caches.get_mut(&spec.id) {
907            let count = entities.len();
908            for (key, entity) in entities {
909                if filter.is_none_or(|filter| filter.matches(&entity)) {
910                    cache.upsert(key, entity);
911                }
912            }
913            // Trim once after the batch; the bound matches the projector's.
914            cache.trim_to_max_entries(max_entries);
915            debug!(view_id = %spec.id, count, "Rebuilt sorted cache from snapshot");
916        }
917    }
918}