ThreadLocal core

Author: slinkydeveloper

sdk-core/src/main/java/dev/restate/sdk/core/sharedcore/SharedCoreImports.java

@@ -8,6 +8,9 @@
 // https://github.com/restatedev/sdk-java/blob/main/LICENSE
 package dev.restate.sdk.core.sharedcore;
 
+import com.dylibso.chicory.annotations.HostModule;
+import com.dylibso.chicory.annotations.WasmExport;
+import com.dylibso.chicory.runtime.HostFunction;
 import com.dylibso.chicory.runtime.ImportValues;
 import com.dylibso.chicory.runtime.Instance;
 import com.dylibso.chicory.runtime.Memory;
@@ -17,22 +20,18 @@
 import dev.restate.sdk.core.sharedcore.generated.SharedCoreWasmMachine;
 import java.io.IOException;
 import java.util.function.Function;
+import org.apache.logging.log4j.Level;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 
 class SharedCoreInstance {
 
-  static final CBORMapper CBOR = CBORMapper.builder()
-          .build();
-  static final int LOG_LEVEL_INFO = 2;
-
-  private static final WasmModule WASM_MODULE;
-
-  static {
-    try {
-      WASM_MODULE = dev.restate.sdk.core.sharedcore.generated.SharedCoreWasm.load();
-    } catch (Exception e) {
-      throw new ExceptionInInitializerError(e);
-    }
-  }
+  private static final Logger LOG = LogManager.getLogger(SharedCoreInstance.class);
+  private static final CBORMapper CBOR = CBORMapper.builder().build();
+  private static final WasmModule WASM_MODULE =
+      dev.restate.sdk.core.sharedcore.generated.SharedCoreWasm.load();
+  private static final ThreadLocal<SharedCoreInstance> THREAD_LOCAL =
+      ThreadLocal.withInitial(SharedCoreInstance::create);
 
   private final Memory memory;
   private final SharedCoreWasm_ModuleExports exports;
@@ -42,22 +41,24 @@ private SharedCoreInstance(Memory memory, SharedCoreWasm_ModuleExports exports)
     this.exports = exports;
   }
 
-  public static SharedCoreInstance create() {
+  static SharedCoreInstance get() {
+    return THREAD_LOCAL.get();
+  }
+
+  private static SharedCoreInstance create() {
     ImportValues importValues =
         ImportValues.builder().addFunction(SharedCoreImports.INSTANCE.toHostFunctions()).build();
 
     Instance instance =
         Instance.builder(WASM_MODULE)
             .withMachineFactory(SharedCoreWasmMachine::new)
-                .withMemoryFactory()
             .withImportValues(importValues)
             .build();
 
     Memory mem = instance.memory();
     SharedCoreWasm_ModuleExports exp = new SharedCoreWasm_ModuleExports(instance);
 
-    // TODO make this configurable
-    exp.init(LOG_LEVEL_INFO);
+    exp.init(toWasmLevel(LOG.getLevel()));
     return new SharedCoreInstance(mem, exp);
   }
 
@@ -136,4 +137,49 @@ public interface QuadFunction<W, X, Y, Z, R> {
   public interface TriConsumer<X, Y, Z> {
     void accept(X x, Y y, Z z);
   }
+
+  static Level toLog4jLevel(int level) {
+    return switch (level) {
+      case 0 -> Level.TRACE;
+      case 1 -> Level.DEBUG;
+      case 2 -> Level.INFO;
+      case 3 -> Level.WARN;
+      default -> Level.ERROR;
+    };
+  }
+
+  static int toWasmLevel(Level level) {
+    if (level == Level.TRACE) {
+      return 0;
+    } else if (level == Level.DEBUG) {
+      return 1;
+    } else if (level == Level.INFO) {
+      return 2;
+    } else if (level == Level.WARN) {
+      return 3;
+    } else {
+      return 4;
+    }
+  }
+
+  @HostModule("env")
+  static final class SharedCoreImports {
+
+    private SharedCoreImports() {}
+
+    public static final SharedCoreImports INSTANCE = new SharedCoreImports();
+
+    @WasmExport
+    public void log(Memory memory, int level, int ptr, int len) {
+      if (len <= 0) {
+        return;
+      }
+      String message = memory.readString(ptr, len);
+      LOG.atLevel(toLog4jLevel(level)).log(message);
+    }
+
+    public HostFunction[] toHostFunctions() {
+      return SharedCoreImports_ModuleFactory.toHostFunctions(this);
+    }
+  }
 }