WIP before removal

Author: slinkydeveloper

.gitignore

@@ -37,4 +37,5 @@ build
 kls_database.db
 .kotlin
 
-.restate
+.restate
+/sdk-core/src/main/rust/target/

gradle/libs.versions.toml

@@ -213,7 +213,32 @@
     [libraries.victools-jsonschema-module-jackson.version]
       ref = 'victools-json-schema'
 
+  [libraries.chicory-runtime]
+    module = 'com.dylibso.chicory:runtime'
+
+    [libraries.chicory-runtime.version]
+      ref = 'chicory'
+
+  [libraries.chicory-annotations]
+    module = 'com.dylibso.chicory:annotations'
+
+    [libraries.chicory-annotations.version]
+      ref = 'chicory'
+
+  [libraries.chicory-annotations-processor]
+    module = 'com.dylibso.chicory:annotations-processor'
+
+    [libraries.chicory-annotations-processor.version]
+      ref = 'chicory'
+
+  [libraries.jackson-cbor]
+    module = 'com.fasterxml.jackson.dataformat:jackson-dataformat-cbor'
+
+    [libraries.jackson-cbor.version]
+      ref = 'jackson'
+
 [plugins]
+  wasm2class = 'at.released.wasm2class.plugin:0.5.0'
   aggregate-javadoc = 'io.freefair.aggregate-javadoc:8.14'
   dependency-license-report = 'com.github.jk1.dependency-license-report:2.9'
   dokka = 'org.jetbrains.dokka:1.9.20'
@@ -233,6 +258,7 @@
       ref = 'ksp'
 
 [versions]
+  chicory = '1.7.5'
   jackson = '2.19.4'
   junit = '5.14.1'
   kotlinx-coroutines = '1.10.2'

sdk-core/build.gradle.kts

@@ -8,16 +8,107 @@ plugins {
   `kotlin-conventions`
   `library-publishing-conventions`
   alias(libs.plugins.jsonschema2pojo)
-  alias(libs.plugins.protobuf)
-  alias(libs.plugins.shadow)
   alias(libs.plugins.ksp)
+  // Chicory AOT: compile .wasm to JVM bytecode at build time
+  id("at.released.wasm2class.plugin") version "0.5.0"
 
   // https://github.com/gradle/gradle/issues/20084#issuecomment-1060822638
   id(libs.plugins.spotless.get().pluginId) apply false
 }
 
 description = "Restate SDK Core"
 
+// ---------------------------------------------------------------------------
+// Rust WASM build pipeline (merged from sdk-shared-core)
+// ---------------------------------------------------------------------------
+
+val rustSrcDir = file("src/main/rust")
+val wasmReleaseDir = file("$rustSrcDir/target/wasm32-unknown-unknown/release")
+val wasmFile = file("$wasmReleaseDir/restate_sdk_shared_core_wasm.wasm")
+val wasmResourceDir = layout.buildDirectory.dir("wasm-resource")
+val wasmResourceFile = wasmResourceDir.map { it.file("restate_sdk_shared_core_wasm.wasm") }
+
+val compileRustToWasm by
+    tasks.registering(Exec::class) {
+      group = "build"
+      description = "Compile the Rust WASM wrapper crate"
+      workingDir = rustSrcDir
+      commandLine("cargo", "build", "--target", "wasm32-unknown-unknown", "--release")
+      inputs.dir("$rustSrcDir/src")
+      inputs.file("$rustSrcDir/Cargo.toml")
+      outputs.file(wasmFile)
+    }
+
+val copyWasm by
+    tasks.registering(Copy::class) {
+      group = "build"
+      dependsOn(compileRustToWasm)
+      from(wasmFile)
+      into(wasmResourceDir)
+    }
+
+// Chicory AOT: compile .wasm → JVM bytecode
+wasm2class {
+  modules {
+    targetPackage = "dev.restate.sdk.core.sharedcore.generated"
+    create("SharedCoreWasm") { wasm = wasmResourceFile }
+  }
+}
+
+tasks.named("precompileWasm2Class") { dependsOn(copyWasm) }
+
+tasks.named("processResources") { dependsOn(copyWasm) }
+
+tasks.withType<KotlinCompile>().configureEach {
+  mustRunAfter(generateWasmMarker, "precompileWasm2Class")
+}
+
+// Generate the @WasmModuleInterface marker class with the absolute wasm file URI so the
+// Chicory annotation processor can find it at compile time (it uses StandardLocation.CLASS_OUTPUT
+// which is not reliable cross-tool, so a file: URI works around that).
+val generatedWasmMarkerDir = layout.buildDirectory.dir("generated-wasm-marker")
+
+val generateWasmMarker by
+    tasks.registering {
+      group = "build"
+      dependsOn(copyWasm)
+      inputs.file(wasmResourceFile)
+      outputs.dir(generatedWasmMarkerDir)
+      doLast {
+        val wasmUri = wasmResourceFile.get().asFile.toURI().toString()
+        val content =
+            """
+// AUTO-GENERATED — do not edit. Regenerated by generateWasmMarker Gradle task.
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+package dev.restate.sdk.core.sharedcore;
+
+import com.dylibso.chicory.annotations.WasmModuleInterface;
+
+/**
+ * Marker class processed by Chicory's annotation processor. Generates
+ * {@code SharedCoreWasm_ModuleExports} (typed wrapper for all WASM exports).
+ */
+@WasmModuleInterface("$wasmUri")
+public final class SharedCoreWasm {}
+    """
+                .trimIndent() + "\n"
+        val out =
+            generatedWasmMarkerDir
+                .get()
+                .file("dev/restate/sdk/core/sharedcore/SharedCoreWasm.java")
+                .asFile
+        out.parentFile.mkdirs()
+        out.writeText(content)
+      }
+    }
+
+// ---------------------------------------------------------------------------
+// Dependency configurations
+// ---------------------------------------------------------------------------
+
 val shade by configurations.creating
 val implementation by configurations.getting
 
@@ -30,17 +121,21 @@ api.extendsFrom(shade)
 dependencies {
   compileOnly(libs.jspecify)
 
-  shadow(project(":sdk-common"))
+  // Chicory annotation processor (@WasmModuleInterface, @HostModule)
+  compileOnly(libs.chicory.annotations)
+  annotationProcessor(libs.chicory.annotations.processor)
 
-  shadow(libs.log4j.api)
-  shadow(libs.opentelemetry.api)
+  implementation(project(":sdk-common"))
+  implementation(libs.log4j.api)
+  implementation(libs.opentelemetry.api)
 
   // We need this for the manifest
-  shadow(libs.jackson.annotations)
-  shadow(libs.jackson.databind)
+  implementation(libs.jackson.annotations)
+  implementation(libs.jackson.databind)
 
-  // We shade protobuf java
-  shade(libs.protobuf.java)
+  // Chicory runtime + Jackson CBOR for the WASM bridge — shaded into the jar
+  implementation(libs.chicory.runtime)
+  implementation(libs.jackson.cbor)
 
   // We don't want a hard-dependency on it
   compileOnly(libs.log4j.core)
@@ -55,10 +150,11 @@ dependencies {
   testImplementation(project(":sdk-api"))
   testImplementation(project(":sdk-api-kotlin"))
   testImplementation(project(":sdk-http-vertx"))
-  testImplementation(project(":sdk-lambda"))
   testImplementation(libs.jackson.annotations)
   testImplementation(libs.jackson.databind)
+  testImplementation(libs.jackson.cbor)
   testImplementation(libs.opentelemetry.api)
+  testImplementation(libs.chicory.runtime)
   testImplementation(libs.protobuf.java)
   testImplementation(libs.mutiny)
   testImplementation(libs.junit.jupiter)
@@ -71,39 +167,42 @@ dependencies {
   testRuntimeOnly(libs.junit.platform.launcher)
 }
 
-// Configure source sets for protobuf plugin and jsonschema2pojo
+// ---------------------------------------------------------------------------
+// Source sets
+// ---------------------------------------------------------------------------
+
 val generatedJ2SPDir = layout.buildDirectory.dir("generated/j2sp")
 
 sourceSets {
   main {
     java.srcDir(generatedJ2SPDir)
-    proto { srcDirs("src/main/service-protocol") }
+    java.srcDir(generatedWasmMarkerDir)
+    proto { srcDirs("src/main/service-protocol") } // TODO(Phase3): remove
+    resources.srcDir(wasmResourceDir)
   }
 }
 
-// Configure jsonSchema2Pojo
+// ---------------------------------------------------------------------------
+// jsonSchema2Pojo
+// ---------------------------------------------------------------------------
+
 jsonSchema2Pojo {
   setSource(files("$projectDir/src/main/service-protocol/endpoint_manifest_schema.json"))
   targetPackage = "dev.restate.sdk.core.generated.manifest"
   targetDirectory = generatedJ2SPDir.get().asFile
-
   useLongIntegers = true
   includeSetters = true
   includeGetters = true
   generateBuilders = true
 }
 
-// Configure protobuf
-
-val protobufVersion = libs.versions.protobuf.get()
-
-protobuf { protoc { artifact = "com.google.protobuf:protoc:$protobufVersion" } }
-
-// Make sure task dependencies are correct
+// ---------------------------------------------------------------------------
+// Task wiring
+// ---------------------------------------------------------------------------
 
 tasks {
   withType<JavaCompile> {
-    dependsOn(generateJsonSchema2Pojo, generateProto)
+    dependsOn(generateJsonSchema2Pojo, generateProto, generateWasmMarker, "precompileWasm2Class")
 
     val disabledClassesCodegen =
         listOf(
@@ -129,25 +228,9 @@ tasks {
   }
   withType<KotlinCompile>().configureEach { dependsOn(generateJsonSchema2Pojo, generateProto) }
   withType<org.gradle.jvm.tasks.Jar>().configureEach {
-    dependsOn(generateJsonSchema2Pojo, generateProto)
+    dependsOn(generateJsonSchema2Pojo, generateProto, generateWasmMarker)
   }
   withType<AbstractDokkaTask>().configureEach { dependsOn(generateJsonSchema2Pojo, generateProto) }
-
-  getByName("jar") {
-    enabled = false
-    dependsOn(shadowJar)
-  }
-
-  shadowJar {
-    configurations = listOf(shade)
-    enableRelocation = true
-    archiveClassifier = null
-    relocate("com.google.protobuf", "dev.restate.shaded.com.google.protobuf")
-    dependencies {
-      project.configurations["shadow"].allDependencies.forEach { exclude(dependency(it)) }
-      exclude("**/google/protobuf/*.proto")
-    }
-  }
 }
 
 ksp {
@@ -169,20 +252,3 @@ ksp {
       )
   arg("dev.restate.codegen.disabledClasses", disabledClassesCodegen.joinToString(","))
 }
-
-// spotless configuration for protobuf
-
-configure<com.diffplug.gradle.spotless.SpotlessExtension> {
-  format("proto") {
-    target("**/*.proto")
-
-    // Exclude proto and service-protocol directories because those get the license header from
-    // their repos.
-    targetExclude(
-        fileTree("$rootDir/sdk-common/src/main/proto") { include("**/*.*") },
-        fileTree("$rootDir/sdk-core/src/main/service-protocol") { include("**/*.*") },
-    )
-
-    licenseHeaderFile("$rootDir/config/license-header", "syntax")
-  }
-}

sdk-core/src/main/java/dev/restate/sdk/core/HandlerContextImpl.java

@@ -441,11 +441,9 @@ private void pollAsyncResultInner(AsyncResultInternal<?> asyncResult) {
 
       if (response instanceof StateMachine.DoProgressResponse.AnyCompleted) {
         // Let it loop now
-      } else if (response instanceof StateMachine.DoProgressResponse.ReadFromInput
-          || response instanceof StateMachine.DoProgressResponse.WaitingPendingRun) {
-        this.stateMachine.onNextEvent(
-            () -> this.pollAsyncResultInner(asyncResult),
-            response instanceof StateMachine.DoProgressResponse.ReadFromInput);
+      } else if (response instanceof StateMachine.DoProgressResponse.WaitExternalProgress
+          || response instanceof StateMachine.DoProgressResponse.CancelSignalReceived) {
+        this.stateMachine.onNextEvent(() -> this.pollAsyncResultInner(asyncResult), true);
         return;
       } else if (response instanceof StateMachine.DoProgressResponse.ExecuteRun) {
         triggerScheduledRun(((StateMachine.DoProgressResponse.ExecuteRun) response).handle());

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

@@ -0,0 +1,47 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// 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.Memory;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+/**
+ * Host module that satisfies the {@code env} imports required by the Rust WASM crate.
+ *
+ * <p>The Rust side imports a single {@code log} function from the {@code env} module and routes all
+ * {@code tracing} output through it. This class bridges those calls to Log4j.
+ */
+@HostModule("env")
+public final class SharedCoreImports {
+
+  private static final Logger LOG = LogManager.getLogger("restate.sdk.shared_core");
+
+  @WasmExport
+  public void log(Memory memory, int level, int ptr, int len) {
+    if (len <= 0) {
+      return;
+    }
+    String message = memory.readString(ptr, len);
+    switch (level) {
+      case 0 -> LOG.trace(message); // TRACE
+      case 1 -> LOG.debug(message); // DEBUG
+      case 2 -> LOG.info(message); // INFO
+      case 3 -> LOG.warn(message); // WARN
+      default -> LOG.error(message); // ERROR (level 4+)
+    }
+  }
+
+  public HostFunction[] toHostFunctions() {
+    return SharedCoreImports_ModuleFactory.toHostFunctions(this);
+  }
+}