Merge pull request #5346 from wasmerio/feat/proc-checkpoint

Implement the WASIX proc_snapshot syscall

lib/cli/src/commands/run/wasi.rs

@@ -172,7 +172,8 @@ pub struct Wasi {
     /// and written to the journal file.
     ///
     /// If not specified, the default is to snapshot when the process idles, when
-    /// the process exits or periodically if an interval argument is also supplied.
+    /// the process exits or periodically if an interval argument is also supplied,
+    /// as well as when the process requests a snapshot explicitly.
     ///
     /// Additionally if the snapshot-on is not specified it will also take a snapshot
     /// on the first stdin, environ or socket listen - this can be used to accelerate

lib/journal/src/concrete/archived.rs

@@ -1604,6 +1604,7 @@ pub enum JournalSnapshotTriggerV1 {
     NonDeterministicCall,
     Bootstrap,
     Transaction,
+    Explicit,
 }
 
 #[repr(C)]

lib/journal/src/concrete/archived_from.rs

@@ -178,6 +178,7 @@ impl From<SnapshotTrigger> for JournalSnapshotTriggerV1 {
             SnapshotTrigger::NonDeterministicCall => JournalSnapshotTriggerV1::NonDeterministicCall,
             SnapshotTrigger::Bootstrap => JournalSnapshotTriggerV1::Bootstrap,
             SnapshotTrigger::Transaction => JournalSnapshotTriggerV1::Transaction,
+            SnapshotTrigger::Explicit => JournalSnapshotTriggerV1::Explicit,
         }
     }
 }
@@ -197,6 +198,7 @@ impl From<JournalSnapshotTriggerV1> for SnapshotTrigger {
             JournalSnapshotTriggerV1::NonDeterministicCall => SnapshotTrigger::NonDeterministicCall,
             JournalSnapshotTriggerV1::Bootstrap => SnapshotTrigger::Bootstrap,
             JournalSnapshotTriggerV1::Transaction => SnapshotTrigger::Transaction,
+            JournalSnapshotTriggerV1::Explicit => SnapshotTrigger::Explicit,
         }
     }
 }
@@ -218,6 +220,7 @@ impl From<&'_ ArchivedJournalSnapshotTriggerV1> for SnapshotTrigger {
             }
             ArchivedJournalSnapshotTriggerV1::Bootstrap => SnapshotTrigger::Bootstrap,
             ArchivedJournalSnapshotTriggerV1::Transaction => SnapshotTrigger::Transaction,
+            ArchivedJournalSnapshotTriggerV1::Explicit => SnapshotTrigger::Explicit,
         }
     }
 }

lib/journal/src/snapshot.rs

@@ -30,22 +30,31 @@ pub enum SnapshotTrigger {
     Bootstrap,
     /// Transaction
     Transaction,
+    /// Explicitly requested by the guest module
+    Explicit,
 }
 
 impl SnapshotTrigger {
     pub fn only_once(&self) -> bool {
         matches!(
             self,
-            Self::FirstListen | Self::FirstEnviron | Self::FirstStdin | Self::FirstSigint
+            Self::FirstListen
+                | Self::FirstEnviron
+                | Self::FirstStdin
+                | Self::FirstSigint
+                // TODO: I don't think this should be an only_once trigger, but
+                // repeatable triggers currently get stuck in a loop
+                | Self::Explicit
         )
     }
 }
 
-pub const DEFAULT_SNAPSHOT_TRIGGERS: [SnapshotTrigger; 4] = [
+pub const DEFAULT_SNAPSHOT_TRIGGERS: [SnapshotTrigger; 5] = [
     SnapshotTrigger::Idle,
     SnapshotTrigger::FirstEnviron,
     SnapshotTrigger::FirstListen,
     SnapshotTrigger::FirstStdin,
+    SnapshotTrigger::Explicit,
 ];
 
 impl FromStr for SnapshotTrigger {
@@ -67,6 +76,7 @@ impl FromStr for SnapshotTrigger {
             "non-deterministic-call" => Self::NonDeterministicCall,
             "bootstrap" => Self::Bootstrap,
             "transaction" => Self::Transaction,
+            "explicit" => Self::Explicit,
             a => return Err(anyhow::format_err!("invalid or unknown trigger ({a})")),
         })
     }

lib/wasix/src/lib.rs

@@ -363,8 +363,7 @@ pub fn generate_import_object_from_env(
 
     let exports_wasi_generic = wasi_exports_generic(store, ctx);
 
-    #[allow(unused_mut)]
-    let mut imports_wasi_generic = imports! {
+    let imports_wasi_generic = imports! {
         "wasi" => exports_wasi_generic,
     };
 
@@ -547,6 +546,7 @@ fn wasix_exports_32(mut store: &mut impl AsStoreMut, env: &FunctionEnv<WasiEnv>)
         "proc_exec2" => Function::new_typed_with_env(&mut store, env, proc_exec2::<Memory32>),
         "proc_raise" => Function::new_typed_with_env(&mut store, env, proc_raise),
         "proc_raise_interval" => Function::new_typed_with_env(&mut store, env, proc_raise_interval),
+        "proc_snapshot" => Function::new_typed_with_env(&mut store, env, proc_snapshot::<Memory32>),
         "proc_spawn" => Function::new_typed_with_env(&mut store, env, proc_spawn::<Memory32>),
         "proc_id" => Function::new_typed_with_env(&mut store, env, proc_id::<Memory32>),
         "proc_parent" => Function::new_typed_with_env(&mut store, env, proc_parent::<Memory32>),
@@ -669,6 +669,7 @@ fn wasix_exports_64(mut store: &mut impl AsStoreMut, env: &FunctionEnv<WasiEnv>)
         "proc_exec2" => Function::new_typed_with_env(&mut store, env, proc_exec2::<Memory64>),
         "proc_raise" => Function::new_typed_with_env(&mut store, env, proc_raise),
         "proc_raise_interval" => Function::new_typed_with_env(&mut store, env, proc_raise_interval),
+        "proc_snapshot" => Function::new_typed_with_env(&mut store, env, proc_snapshot::<Memory64>),
         "proc_spawn" => Function::new_typed_with_env(&mut store, env, proc_spawn::<Memory64>),
         "proc_id" => Function::new_typed_with_env(&mut store, env, proc_id::<Memory64>),
         "proc_parent" => Function::new_typed_with_env(&mut store, env, proc_parent::<Memory64>),

lib/wasix/src/syscalls/wasix/mod.rs

@@ -28,6 +28,7 @@ mod proc_id;
 mod proc_join;
 mod proc_parent;
 mod proc_signal;
+mod proc_snapshot;
 mod proc_spawn;
 mod resolve;
 mod sched_yield;
@@ -97,6 +98,7 @@ pub use proc_id::*;
 pub use proc_join::*;
 pub use proc_parent::*;
 pub use proc_signal::*;
+pub use proc_snapshot::*;
 pub use proc_spawn::*;
 pub use resolve::*;
 pub use sched_yield::*;

lib/wasix/src/syscalls/wasix/proc_snapshot.rs

@@ -0,0 +1,11 @@
+use super::*;
+use crate::syscalls::*;
+
+/// Explicitly snapshots the process state.
+#[instrument(level = "trace", skip_all, ret)]
+pub fn proc_snapshot<M: MemorySize>(
+    mut ctx: FunctionEnvMut<'_, WasiEnv>,
+) -> Result<Errno, WasiError> {
+    wasi_try_ok!(maybe_snapshot_once::<M>(ctx, SnapshotTrigger::Explicit)?);
+    Ok(Errno::Success)
+}