Function test_custom_memory_limiter_async

tests/all/limits.rs:544–622

()

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_custom_memory_limiter_async() -> Result<()> {
    let engine = Engine::default();
    let mut linker = Linker::new(&engine);

    // This approximates a function that would "allocate" resources that the host tracks.
    // Here this is a simple function that increments the current host memory "used".
    linker.func_wrap(
        "",
        "alloc",
        |mut caller: Caller<'_, MemoryContext>, size: u32| -> u32 {
            let ctx = caller.data_mut();
            let size = size as usize;

            if size + ctx.host_memory_used + ctx.wasm_memory_used <= ctx.memory_limit {
                ctx.host_memory_used += size;
                return 1;
            }

            ctx.limit_exceeded = true;

            0
        },
    )?;

    let module = Module::new(
        &engine,
        r#"(module (import "" "alloc" (func $alloc (param i32) (result i32))) (memory (export "m") 0) (func (export "f") (param i32) (result i32) local.get 0 call $alloc))"#,
    )?;

    let context = MemoryContext {
        host_memory_used: 0,
        wasm_memory_used: 0,
        memory_limit: 1 << 20, // 16 wasm pages is the limit for both wasm + host memory
        limit_exceeded: false,
    };

    let mut store = Store::new(&engine, context);
    store.limiter_async(|s| s as &mut dyn ResourceLimiterAsync);
    let instance = linker.instantiate_async(&mut store, &module).await?;
    let memory = instance.get_memory(&mut store, "m").unwrap();

    // Grow the memory by 640 KiB
    memory.grow_async(&mut store, 3).await?;
    memory.grow_async(&mut store, 5).await?;
    memory.grow_async(&mut store, 2).await?;

    assert!(!store.data().limit_exceeded);

    // Grow the host "memory" by 384 KiB
    let f = instance.get_typed_func::<u32, u32>(&mut store, "f")?;

    assert_eq!(f.call_async(&mut store, 1 * 0x10000).await?, 1);
    assert_eq!(f.call_async(&mut store, 3 * 0x10000).await?, 1);
    assert_eq!(f.call_async(&mut store, 2 * 0x10000).await?, 1);

    // Memory is at the maximum, but the limit hasn't been exceeded
    assert!(!store.data().limit_exceeded);