Allow configuring a machine factory on plugins.

proxy-wasm-java-host/src/main/java/io/roastedroot/proxywasm/ProxyWasm.java

@@ -2,11 +2,13 @@
 
 import static io.roastedroot.proxywasm.Helpers.len;
 
+import com.dylibso.chicory.experimental.aot.AotMachine;
 import com.dylibso.chicory.runtime.ByteArrayMemory;
 import com.dylibso.chicory.runtime.HostFunction;
 import com.dylibso.chicory.runtime.ImportMemory;
 import com.dylibso.chicory.runtime.ImportValues;
 import com.dylibso.chicory.runtime.Instance;
+import com.dylibso.chicory.runtime.Machine;
 import com.dylibso.chicory.wasi.WasiOptions;
 import com.dylibso.chicory.wasi.WasiPreview1;
 import com.dylibso.chicory.wasm.WasmModule;
@@ -18,6 +20,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
 
 public final class ProxyWasm implements Closeable {
 
@@ -226,6 +229,7 @@ public static class Builder implements Cloneable {
         private ImportMemory memory;
         private WasiOptions wasiOptions;
         private boolean start = true;
+        private Function<Instance, Machine> machineFactory;
 
         @Override
         @SuppressWarnings("NoClone")
@@ -261,6 +265,11 @@ public ProxyWasm.Builder withWasiOptions(WasiOptions options) {
             return this;
         }
 
+        public ProxyWasm.Builder withMachineFactory(Function<Instance, Machine> machineFactory) {
+            this.machineFactory = machineFactory;
+            return this;
+        }
+
         Builder() {}
 
         public ProxyWasm build(Instance instance) throws StartException {
@@ -275,6 +284,10 @@ public ProxyWasm build(WasmModule module) throws StartException {
         public ProxyWasm build(Instance.Builder instanceBuilder) throws StartException {
             var imports = ImportValues.builder();
 
+            if (this.machineFactory != null) {
+                instanceBuilder.withMachineFactory(AotMachine::new);
+            }
+
             imports.addMemory(Objects.requireNonNullElseGet(memory, this::defaultImportMemory));
             imports.addFunction(toHostFunctions());
             imports.addFunction(

proxy-wasm-java-host/src/main/java/io/roastedroot/proxywasm/plugin/Plugin.java

@@ -10,6 +10,8 @@
 
 import com.dylibso.chicory.runtime.ImportMemory;
 import com.dylibso.chicory.runtime.Instance;
+import com.dylibso.chicory.runtime.Machine;
+import com.dylibso.chicory.wasi.WasiOptions;
 import com.dylibso.chicory.wasm.WasmModule;
 import io.roastedroot.proxywasm.ArrayProxyMap;
 import io.roastedroot.proxywasm.ChainedHandler;
@@ -35,6 +37,7 @@
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Function;
 
 public final class Plugin {
 
@@ -187,36 +190,46 @@ public Builder withSharedDataHandler(SharedDataHandler sharedDataHandler) {
             return this;
         }
 
-        public Plugin.Builder withShared(boolean shared) {
+        public Builder withShared(boolean shared) {
             this.shared = shared;
             return this;
         }
 
-        public Plugin.Builder withVmConfig(byte[] vmConfig) {
+        public Builder withVmConfig(byte[] vmConfig) {
             this.vmConfig = vmConfig;
             return this;
         }
 
-        public Plugin.Builder withVmConfig(String vmConfig) {
+        public Builder withVmConfig(String vmConfig) {
             this.vmConfig = bytes(vmConfig);
             return this;
         }
 
-        public Plugin.Builder withPluginConfig(byte[] pluginConfig) {
+        public Builder withPluginConfig(byte[] pluginConfig) {
             this.pluginConfig = pluginConfig;
             return this;
         }
 
-        public Plugin.Builder withPluginConfig(String pluginConfig) {
+        public Builder withPluginConfig(String pluginConfig) {
             this.pluginConfig = bytes(pluginConfig);
             return this;
         }
 
-        public Plugin.Builder withImportMemory(ImportMemory memory) {
+        public Builder withImportMemory(ImportMemory memory) {
             proxyWasmBuilder = proxyWasmBuilder.withImportMemory(memory);
             return this;
         }
 
+        public Builder withMachineFactory(Function<Instance, Machine> machineFactory) {
+            proxyWasmBuilder.withMachineFactory(machineFactory);
+            return this;
+        }
+
+        public Builder withWasiOptions(WasiOptions options) {
+            proxyWasmBuilder.withWasiOptions(options);
+            return this;
+        }
+
         public Plugin build(WasmModule module) throws StartException {
             return build(proxyWasmBuilder.build(module));
         }

proxy-wasm-jaxrs-jersey/src/test/java/io/roastedroot/proxywasm/jaxrs/example/App.java

@@ -1,5 +1,6 @@
 package io.roastedroot.proxywasm.jaxrs.example;
 
+import com.dylibso.chicory.experimental.aot.AotMachine;
 import com.dylibso.chicory.wasm.Parser;
 import com.dylibso.chicory.wasm.WasmModule;
 import com.google.gson.Gson;
@@ -24,6 +25,7 @@ public static PluginFactory headerTests() throws StartException {
                         .withName("headerTests")
                         .withLogger(new MockLogger("headerTests"))
                         .withPluginConfig(gson.toJson(Map.of("type", "headerTests")))
+                        .withMachineFactory(AotMachine::new)
                         .build(parseTestModule("/go-examples/unit_tester/main.wasm"));
     }
 
@@ -34,6 +36,7 @@ public static PluginFactory headerTestsNotShared() throws StartException {
                         .withShared(false)
                         .withLogger(new MockLogger("headerTestsNotShared"))
                         .withPluginConfig(gson.toJson(Map.of("type", "headerTests")))
+                        .withMachineFactory(AotMachine::new)
                         .build(parseTestModule("/go-examples/unit_tester/main.wasm"));
     }
 
@@ -43,6 +46,7 @@ public static PluginFactory tickTests() throws StartException {
                         .withName("tickTests")
                         .withLogger(new MockLogger("tickTests"))
                         .withPluginConfig(gson.toJson(Map.of("type", "tickTests")))
+                        .withMachineFactory(AotMachine::new)
                         .build(parseTestModule("/go-examples/unit_tester/main.wasm"));
     }
 
@@ -53,6 +57,7 @@ public static PluginFactory ffiTests() throws StartException {
                         .withLogger(new MockLogger("ffiTests"))
                         .withPluginConfig(gson.toJson(Map.of("type", "ffiTests")))
                         .withForeignFunctions(Map.of("reverse", App::reverse))
+                        .withMachineFactory(AotMachine::new)
                         .build(parseTestModule("/go-examples/unit_tester/main.wasm"));
     }
 
@@ -76,6 +81,7 @@ public static PluginFactory httpCallTests() throws StartException {
                                                 "upstream", "web_service",
                                                 "path", "/ok")))
                         .withUpstreams(Map.of("web_service", "localhost:8081"))
+                        .withMachineFactory(AotMachine::new)
                         .build(parseTestModule("/go-examples/unit_tester/main.wasm"));
     }
 }