clusterflux-public/crates/disasmer-node/src/bin/disasmer-wasmtime-smoke.rs

use std::fs;
use std::path::PathBuf;

use disasmer_core::{CommandInvocation, NodeId, TaskId, VfsPath};
use disasmer_node::{LocalCommandExecutor, VirtualThreadCommand, WasmtimeTaskRuntime};
use serde_json::json;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() == 2 && args[1] == "--host-command" {
        return run_host_command_smoke();
    }
    if args.len() == 6 && args[1] == "--debug-freeze-resume" {
        return run_debug_freeze_resume_smoke(&args[2], &args[3], &args[4], &args[5]);
    }
    if args.len() != 5 {
        return Err(
            "usage: disasmer-wasmtime-smoke <module.wasm> <export> <i32-arg> <expected-i32> | --host-command | --debug-freeze-resume <module.wasm> <export> <i32-arg> <expected-i32>".into(),
        );
    }

    let module = PathBuf::from(&args[1]);
    let export = &args[2];
    let arg: i32 = args[3].parse()?;
    let expected: i32 = args[4].parse()?;

    let wasm = fs::read(&module)?;
    let runtime = WasmtimeTaskRuntime::new()?;
    let result = runtime.run_i32_export(&wasm, export, arg)?;
    if result != expected {
        return Err(format!("expected {expected}, got {result} from export `{export}`").into());
    }

    println!(
        "{}",
        serde_json::to_string(&json!({
            "type": "wasmtime_task_smoke",
            "module": module,
            "export": export,
            "arg": arg,
            "result": result,
        }))?
    );
    Ok(())
}

fn run_debug_freeze_resume_smoke(
    module: &str,
    export: &str,
    arg: &str,
    expected: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let module = PathBuf::from(module);
    let arg: i32 = arg.parse()?;
    let expected: i32 = expected.parse()?;
    let wasm = fs::read(&module)?;
    let runtime = WasmtimeTaskRuntime::new()?;
    let probe = runtime.freeze_resume_i32_export_probe(&wasm, export, arg)?;
    if probe.result != expected {
        return Err(format!(
            "expected {expected}, got {} from export `{export}`",
            probe.result
        )
        .into());
    }

    println!(
        "{}",
        serde_json::to_string(&json!({
            "type": "wasmtime_debug_freeze_resume_smoke",
            "module": module,
            "export": export,
            "task": probe.task,
            "frozen_state": probe.frozen_state,
            "resumed_state": probe.resumed_state,
            "stack_frames": probe.stack_frames,
            "arg": arg,
            "result": probe.result,
            "node_runtime_reached_wasm_task": true,
        }))?
    );
    Ok(())
}

fn run_host_command_smoke() -> Result<(), Box<dyn std::error::Error>> {
    let runtime = WasmtimeTaskRuntime::new()?;
    let result = runtime.run_i32_export_with_command_import(
        r#"
        (module
          (import "disasmer" "cmd_run" (func $cmd_run (result i32)))
          (func (export "compile-linux") (result i32)
            call $cmd_run))
        "#,
        "compile-linux",
        LocalCommandExecutor {
            node: NodeId::from("node-wasmtime"),
            hosted_control_plane: false,
            has_command_capability: true,
        },
        VirtualThreadCommand {
            virtual_thread: TaskId::from("compile-linux"),
            invocation: CommandInvocation {
                program: "sh".to_owned(),
                args: vec!["-c".to_owned(), "printf linux-build-artifact".to_owned()],
                env: None,
            },
            stage_stdout_as: Some(VfsPath::new("/vfs/artifacts/linux/app.tar.zst")?),
        },
    )?;

    println!(
        "{}",
        serde_json::to_string(&json!({
            "type": "wasmtime_host_command_smoke",
            "export": "compile-linux",
            "export_result": result.export_result,
            "virtual_thread": result.command_output.virtual_thread,
            "stdout": result.command_output.stdout,
            "staged_artifact": result.command_output.staged_artifact,
            "large_bytes_uploaded": result.manifest.large_bytes_uploaded,
            "manifest_objects": result.manifest.objects.len(),
            "node_host_import": "disasmer.cmd_run",
            "flagship_linux_build_task": true,
            "node_executed_host_command": true,
            "hosted_control_plane_ran_command": false,
        }))?
    );
    Ok(())
}