clusterflux-public/crates/clusterflux-wasm-runtime/src/tests.rs

use super::*;

#[test]
fn fuel_token_bucket_refills_after_idle_and_never_exceeds_burst() {
    let limits = WasmtimeRuntimeLimits {
        fuel_units_per_second: 100,
        fuel_burst_seconds: 2,
        memory_bytes: 1024 * 1024,
    };
    let mut bucket = FuelTokenBucket::new(&limits);
    bucket.last_refill = Instant::now() - Duration::from_millis(1_500);
    let refilled = bucket.refill(0);
    assert!((150..=151).contains(&refilled));

    bucket.last_refill = Instant::now() - Duration::from_secs(10);
    assert_eq!(bucket.refill(refilled), 200);
    for _ in 0..10_000 {
        assert!(bucket.refill(200) <= 200);
    }
}

struct AbortSignalHost {
    abort: Arc<AtomicBool>,
    debug: Option<Arc<WasmDebugControl>>,
}

impl WasmTaskHost for AbortSignalHost {
    fn abort_signal(&self) -> Option<Arc<AtomicBool>> {
        Some(Arc::clone(&self.abort))
    }

    fn debug_control(&self) -> Option<Arc<WasmDebugControl>> {
        self.debug.clone()
    }

    fn start_task(
        &mut self,
        _request: WasmHostTaskStartRequest,
    ) -> Result<WasmHostTaskHandle, String> {
        Err("not used".to_owned())
    }

    fn join_task(
        &mut self,
        _request: WasmHostTaskJoinRequest,
    ) -> Result<WasmHostTaskJoinResult, String> {
        Err("not used".to_owned())
    }

    fn run_command(
        &mut self,
        _request: WasmHostCommandRequest,
    ) -> Result<WasmHostCommandResult, String> {
        Err("not used".to_owned())
    }

    fn poll_task_control(
        &mut self,
        request: WasmHostTaskControlRequest,
    ) -> Result<WasmHostTaskControlResult, String> {
        request.validate()?;
        Ok(WasmHostTaskControlResult {
            abi_version: clusterflux_core::WASM_TASK_ABI_VERSION,
            cancellation_requested: false,
        })
    }

    fn vfs_operation(&mut self, _request: WasmHostVfsRequest) -> Result<WasmHostVfsResult, String> {
        Err("not used".to_owned())
    }

    fn snapshot_source(
        &mut self,
        _request: WasmHostSourceSnapshotRequest,
    ) -> Result<WasmHostSourceSnapshotResult, String> {
        Err("not used".to_owned())
    }
}

#[test]
fn epoch_interruption_aborts_cpu_bound_wasm_without_a_host_call() {
    let abort = Arc::new(AtomicBool::new(false));
    let trigger = Arc::clone(&abort);
    let trigger_thread = thread::spawn(move || {
        thread::sleep(Duration::from_millis(50));
        trigger.store(true, Ordering::Release);
    });
    let wasm = r#"
            (module
              (func (export "spin") (param i32) (result i32)
                (loop $forever
                  br $forever)
                i32.const 0))
        "#;
    let started = std::time::Instant::now();
    let error = WasmtimeTaskRuntime::new()
        .unwrap()
        .run_i32_export_verified_with_task_host(
            wasm,
            &Digest::sha256(wasm),
            "spin",
            1,
            Box::new(AbortSignalHost { abort, debug: None }),
        )
        .unwrap_err();
    trigger_thread.join().unwrap();

    assert!(matches!(error, WasmTaskError::HostControl(_)));
    assert!(error.to_string().contains("process abort"));
    assert!(started.elapsed() < Duration::from_secs(2));
}

#[test]
fn aborting_one_store_does_not_poison_the_shared_engine() {
    let runtime = WasmtimeTaskRuntime::new().unwrap();
    let abort = Arc::new(AtomicBool::new(false));
    let trigger = Arc::clone(&abort);
    let trigger_thread = thread::spawn(move || {
        thread::sleep(Duration::from_millis(50));
        trigger.store(true, Ordering::Release);
    });
    let spinning_wasm = r#"
            (module
              (func (export "spin") (param i32) (result i32)
                (loop $forever
                  br $forever)
                i32.const 0))
        "#;
    let error = runtime
        .run_i32_export_verified_with_task_host(
            spinning_wasm,
            &Digest::sha256(spinning_wasm),
            "spin",
            0,
            Box::new(AbortSignalHost { abort, debug: None }),
        )
        .unwrap_err();
    trigger_thread.join().unwrap();
    assert!(matches!(error, WasmTaskError::HostControl(_)));

    let succeeding_wasm = r#"
            (module
              (func (export "increment") (param i32) (result i32)
                local.get 0
                i32.const 1
                i32.add))
        "#;
    assert_eq!(
        runtime
            .run_i32_export_verified(
                succeeding_wasm,
                &Digest::sha256(succeeding_wasm),
                "increment",
                41,
            )
            .unwrap(),
        42
    );
}

#[test]
fn debug_control_freezes_and_resumes_executing_wasm_before_abort() {
    let runtime = WasmtimeTaskRuntime::new().unwrap();
    let abort = Arc::new(AtomicBool::new(false));
    let debug = Arc::new(WasmDebugControl::default());
    let execution_abort = Arc::clone(&abort);
    let execution_debug = Arc::clone(&debug);
    let (finished_tx, finished_rx) = std::sync::mpsc::channel();
    let execution = thread::spawn(move || {
        let wasm = r#"
                (module
                  (func (export "spin") (param i32) (result i32)
                    (loop $forever
                      br $forever)
                    i32.const 0))
            "#;
        let result = runtime.run_i32_export_verified_with_task_host(
            wasm,
            &Digest::sha256(wasm),
            "spin",
            0,
            Box::new(AbortSignalHost {
                abort: execution_abort,
                debug: Some(execution_debug),
            }),
        );
        let _ = finished_tx.send(());
        result
    });

    debug.request_freeze(7);
    assert!(debug.wait_until_frozen(7, Duration::from_secs(2)));
    assert!(finished_rx.try_recv().is_err());
    debug.request_resume(7);
    assert!(debug.wait_until_running(7, Duration::from_secs(2)));
    abort.store(true, Ordering::Release);

    let error = execution.join().unwrap().unwrap_err();
    assert!(matches!(error, WasmTaskError::HostControl(_)));
    assert!(error.to_string().contains("process abort"));
}

#[test]
fn pending_freeze_blocks_a_quiescent_host_call_before_it_starts() {
    let debug = Arc::new(WasmDebugControl::default());
    debug.request_freeze(9);
    let boundary_debug = Arc::clone(&debug);
    let (entered_tx, entered_rx) = std::sync::mpsc::channel();
    let boundary = thread::spawn(move || {
        boundary_debug.enter_quiescent_host_boundary(None);
        entered_tx.send(()).unwrap();
        boundary_debug.leave_quiescent_host_boundary(None);
    });

    assert!(debug.wait_until_frozen(9, Duration::from_secs(1)));
    assert!(entered_rx.try_recv().is_err());
    debug.request_resume(9);
    assert!(debug.wait_until_running(9, Duration::from_secs(1)));
    entered_rx.recv_timeout(Duration::from_secs(1)).unwrap();
    boundary.join().unwrap();
}

#[test]
fn wasm_linear_memory_growth_is_bounded_per_store() {
    let wasm = r#"
        (module
          (memory 1)
          (func (export "grow") (param i32) (result i32)
            i32.const 5000
            memory.grow
            drop
            local.get 0))
    "#;
    let runtime = WasmtimeTaskRuntime::new().unwrap();
    let error = runtime.run_i32_export(wasm, "grow", 7).unwrap_err();

    assert!(error.to_string().contains("memory") || error.to_string().contains("grow"));
    assert_eq!(
        runtime
            .run_i32_export(
                "(module (func (export \"healthy\") (param i32) (result i32) local.get 0))",
                "healthy",
                11,
            )
            .unwrap(),
        11
    );
}