Memory access in Wasi module built with TinyGo

[nicholasjackson]

Hi,

I have been trying out the memory example and while I can get this to work with Rust and AssemblyScript compiled WASM modules I can't seem to get this working with Go code compile with WASI.

When I try to access the memory using the following code

    Optional<Extern> ext = linker.get(store, "", "memory");
    Memory mem = ext.get().memory();

ext is always null from the Go Wasm module.

I have compiled the module using the following flags which have worked for me when I was using the wasmer library.

    tinygo build -o ../go.wasm -wasm-abi=generic -target=wasi -scheduler=coroutines -gc=leaking .

When I look at the wat file (attached inline), I can see the memory exported however for some reason this can not be found by wasmtime.

 (table (;0;) 1 1 funcref)
  (memory (;0;) 2)
  (global (;0;) (mut i32) (i32.const 65536))
  (export "memory" (memory 0))
  (export "_start" (func $_start))
  (export "allocate" (func $allocate))
  (export "deallocate" (func $deallocate))
  (export "get_string_size" (func $get_string_size))
  (export "sum" (func $sum))
  (export "hello" (func $hello))

I have also included my example code that I am using run your library. Great library by the way, I am giving a talk at a conference next week and I am looking forward to demoing this.

Example Wat: go.zip

Original Source

package com.github.nicholasjackson.wasmcraft.wasm;

import static io.github.kawamuray.wasmtime.WasmValType.I32;

import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.MessageDigest;
import java.util.ArrayList;
import java.util.Hashtable;
import java.util.Optional;

import io.github.kawamuray.wasmtime.Engine;
import io.github.kawamuray.wasmtime.Extern;
import io.github.kawamuray.wasmtime.Func;
import io.github.kawamuray.wasmtime.Linker;
import io.github.kawamuray.wasmtime.Memory;
import io.github.kawamuray.wasmtime.Module;
import io.github.kawamuray.wasmtime.Store;
import io.github.kawamuray.wasmtime.WasmFunctions;
import io.github.kawamuray.wasmtime.wasi.WasiCtx;
import io.github.kawamuray.wasmtime.wasi.WasiCtxBuilder;

public class WasmRuntime {

  private Store<Void> store;
  private Engine engine;

  private Hashtable<String, Module> modules = new Hashtable<>();

  private static WasmRuntime instance;

  public WasmRuntime() {
  }

  public static WasmRuntime getInstance() {
    if (instance == null) {
      instance = new WasmRuntime();
    }

    return instance;
  }

  public void init() {
    WasiCtx wasi = new WasiCtxBuilder().inheritStdout().inheritStderr().build();

    store = Store.withoutData(wasi);
    engine = store.engine();
  }

  // getModule returns the module at the given path.
  // If the module has not been loaded, it will be loaded and cached.
  public String getModule(String uri) throws Exception {
    // C:\Users\jacks\java\fabric-wasm-mod\wasm\sum.wat
    System.out.println("Compiling module: " + uri);

    // get the hash for the file
    String hash = computeHashForFile(uri);

    // check if the module has been loaded
    if (modules.containsKey(hash)) {
      return hash;
    }

    // module has not been loaded, load it
    Module module = Module.fromFile(engine, uri);

    // cache the module
    modules.put(hash, module);

    return hash;
  }

  public Object executeModuleFunction(String moduleHash, String functionName, String... args) throws Exception {
    Module module = modules.get(moduleHash);

    Linker linker = new Linker(engine);
    WasiCtx.addToLinker(linker);
    linker.module(store, "", module);

    // create a list of parameters, it is possible that this list will be empty
    ArrayList<Integer> params = new ArrayList<Integer>();

    for (String arg : args) {
      if (arg.isBlank()) {
        continue;
      }

      Integer param = 0;

      try {
        // is the argument an integer
        param = Integer.parseInt(arg);
      } catch (Exception e) {
        // not an integer so create the string in the wasm modules memory and pass a
        // reference to that string as an integer
        param = setStringInMemory(arg, linker);
      }

      // add to the collection
      params.add(param);
    }

    Integer result = 0;

    // get an instance of the function from the module
    Func f = linker.get(store, "", functionName).get().func();

    switch (params.size()) {
      case 0:
        WasmFunctions.Function0<Integer> fn0 = WasmFunctions.func(store, f, I32);
        result = fn0.call();
        break;
      case 1:
        WasmFunctions.Function1<Integer, Integer> fn1 = WasmFunctions.func(store, f, I32, I32);
        result = fn1.call(params.get(0));
        break;
      case 2:
        WasmFunctions.Function2<Integer, Integer, Integer> fn2 = WasmFunctions.func(store, f, I32, I32, I32);
        result = fn2.call(params.get(0), params.get(0));
        break;
    }

    // close and free up resources
    linker.close();

    return result;
  }

  private String computeHashForFile(String uri) {
    try {
      byte[] b = Files.readAllBytes(Paths.get(uri));
      byte[] hash = MessageDigest.getInstance("MD5").digest(b);
      return new String(hash);
    } finally {
      return "";
    }
  }

  private Integer setStringInMemory(String string, Linker linker) {
    // allocate the memory
    Func f = linker.get(store, "", "allocate").get().func();

    WasmFunctions.Function1<Integer, Integer> fn = WasmFunctions.func(store, f, I32, I32);
    Integer result = fn.call(string.length() + 1);

    Optional<Extern> ext = linker.get(store, "", "memory");
    Memory mem = ext.get().memory();

    ByteBuffer buf = mem.buffer(store);

    for (int i = 0; i < string.length(); i++) {
      buf.put(result + i, (byte) string.charAt(i));
    }

    // add the null terminator
    byte nullTerminator = 0;
    buf.put(result + string.length(), nullTerminator);

    return result;
  }
}

[nicholasjackson]

I think I might have found the answer as you detailed in this issue: https://github.com/kawamuray/wasmtime-java/issues/10#issuecomment-762309408

Now to figure how to configure as a reactor-type

[nicholasjackson]

Far from ideal but this does work, I will raise an issue in TinyGo to see if there is a better way round this.

    tinygo build -o ../go.wasm -wasm-abi=generic -target=wasi -scheduler=coroutines -gc=leaking .
    wasm2wat ../go.wasm > ../go.wat

# Remove _start line from go.wat
    sed -i '/.*export "_start".*/d' ../go.wat
    wat2wasm ../go.wat --output ../go.wasm

[kawamuray]

Hi, thanks for reporting the issue.

As you already figured out, a memory instance isn't ready yet at the time you try to get export because your module is recognized as a command, which creates an instance only after a function of the module is called, and destroys it before returns.

So for what you'd like to demonstrate, generating a wasm module as a reactor is right approach, but accessing memory instance is still possible with command type module as long as it happens within the callstack of the first function, such as in callback provided by host (java) side.

Here's an example, assuming your module's "main" function calls an imported function "xyz::delegate_main",

// Create a Func by yourself with low level interface, so that your handler has access to Caller, which provides getExport
Func delegateMain = new Func(store,  FuncType.empty(), (caller, params, results) -> {
    Memory mem = caller.getExport("memory").get().memory();
    mem... // play with it
    return Optional.empty();
});
linker.define("xyz", "delegate_main", Extern.fromFunc(delegateMain));
invokeMain(linker.get(store, "", "main").get().func()); // An exported memory will be touched within the callstack of this call

Great library by the way, I am giving a talk at a conference next week and I am looking forward to demoing this.

Thanks! I'm happy to hear that :)

[nicholasjackson]

Awesome thanks, I have another question on files but will start a different thread.