A bunch of fixes for the Workflow API (#222)

Author: slinkydeveloper

sdk-api-gen/src/main/resources/templates/ServiceAdapter.hbs

@@ -2,10 +2,12 @@
 
 public class {{className}}ServiceAdapter implements dev.restate.sdk.common.ServiceAdapter<{{fqcn}}> {
 
+    public static final String SERVICE_NAME = "{{fqcn}}";
+
     @java.lang.Override
     public dev.restate.sdk.workflow.impl.WorkflowServicesBundle adapt({{fqcn}} service) {
         return dev.restate.sdk.workflow.impl.WorkflowServicesBundle.named(
-            "{{fqcn}}",
+            SERVICE_NAME,
             {{#methods}}{{#if isWorkflow}}
             dev.restate.sdk.workflow.impl.WorkflowServicesBundle.MethodSignature.of("{{name}}", {{{inputSerdeDecl}}}, {{{outputSerdeDecl}}}),
             (ctx, req) -> {

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

@@ -10,7 +10,6 @@
 
 import dev.restate.generated.service.discovery.Discovery;
 import dev.restate.sdk.common.ServiceAdapter;
-import dev.restate.sdk.common.ServicesBundle;
 import io.grpc.*;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.trace.Span;
@@ -221,11 +220,12 @@ public void setInvocationStatus(String invocationStatus) {
   }
 
   /** Resolve the code generated {@link ServiceAdapter} */
-  public static ServicesBundle adapt(Object entity) {
+  public static ServiceAdapter<Object> discoverAdapter(Object entity) {
     Class<?> userClazz = entity.getClass();
 
     // Find Service code-generated class
     // TODO This could be done with an SPI
+    // TODO This should support interfaces
     Class<?> serviceAdapterClazz;
     try {
       serviceAdapterClazz = Class.forName(userClazz.getCanonicalName() + "ServiceAdapter");
@@ -251,6 +251,6 @@ public static ServicesBundle adapt(Object entity) {
           e);
     }
 
-    return serviceAdapter.adapt(entity);
+    return serviceAdapter;
   }
 }

sdk-http-vertx/src/main/java/dev/restate/sdk/http/vertx/RestateHttpEndpointBuilder.java

@@ -11,6 +11,7 @@
 import dev.restate.generated.service.discovery.Discovery;
 import dev.restate.sdk.common.BlockingService;
 import dev.restate.sdk.common.NonBlockingService;
+import dev.restate.sdk.common.ServiceAdapter;
 import dev.restate.sdk.core.RestateEndpoint;
 import io.grpc.ServerInterceptor;
 import io.grpc.ServerInterceptors;
@@ -117,20 +118,42 @@ public RestateHttpEndpointBuilder withService(
     return this;
   }
 
-  /** Add a Restate entity to the endpoint. */
+  /**
+   * Add a Restate service to the endpoint. This will automatically discover the adapter based on
+   * the class name. You can provide the adapter manually using {@link #with(Object,
+   * ServiceAdapter)}
+   */
   public RestateHttpEndpointBuilder with(Object service) {
     return this.with(service, defaultExecutor);
   }
 
   /**
-   * Add a Restate entity to the endpoint, specifying the {@code executor} where to run the entity
-   * code.
+   * Add a Restate service to the endpoint, specifying the {@code executor} where to run the service
+   * code. This will automatically discover the adapter based on the class name. You can provide the
+   * adapter manually using {@link #with(Object, ServiceAdapter, Executor)}
    *
    * <p>You can run on virtual threads by using the executor {@code
    * Executors.newVirtualThreadPerTaskExecutor()}.
    */
   public RestateHttpEndpointBuilder with(Object service, Executor executor) {
-    List<BlockingService> services = RestateEndpoint.adapt(service).services();
+    return this.with(service, RestateEndpoint.discoverAdapter(service), executor);
+  }
+
+  /** Add a Restate service to the endpoint, specifying an adapter. */
+  public <T> RestateHttpEndpointBuilder with(T service, ServiceAdapter<T> adapter) {
+    return this.with(service, adapter, defaultExecutor);
+  }
+
+  /**
+   * Add a Restate service to the endpoint, specifying the {@code executor} where to run the service
+   * code.
+   *
+   * <p>You can run on virtual threads by using the executor {@code
+   * Executors.newVirtualThreadPerTaskExecutor()}.
+   */
+  public <T> RestateHttpEndpointBuilder with(
+      T service, ServiceAdapter<T> adapter, Executor executor) {
+    List<BlockingService> services = adapter.adapt(service).services();
     for (BlockingService svc : services) {
       this.withService(svc, executor);
     }

sdk-lambda/src/main/java/dev/restate/sdk/lambda/RestateLambdaEndpointBuilder.java

@@ -11,6 +11,7 @@
 import dev.restate.generated.service.discovery.Discovery;
 import dev.restate.sdk.common.BlockingService;
 import dev.restate.sdk.common.Service;
+import dev.restate.sdk.common.ServiceAdapter;
 import dev.restate.sdk.core.RestateEndpoint;
 import io.grpc.ServerInterceptor;
 import io.grpc.ServerInterceptors;
@@ -44,8 +45,12 @@ public RestateLambdaEndpointBuilder withService(
    * code.
    */
   public RestateLambdaEndpointBuilder with(Object service) {
-    List<BlockingService> services = RestateEndpoint.adapt(service).services();
-    for (BlockingService svc : services) {
+    return this.with(service, RestateEndpoint.discoverAdapter(service));
+  }
+
+  public <T> RestateLambdaEndpointBuilder with(T service, ServiceAdapter<T> adapter) {
+    List<BlockingService> services = adapter.adapt(service).services();
+    for (Service svc : services) {
       this.restateGrpcServerBuilder.withService(svc);
     }
 

sdk-workflow-api/src/main/java/dev/restate/sdk/workflow/impl/WorkflowManagerImpl.java

@@ -135,10 +135,10 @@ public StartResponse tryStart(KeyedContext context, StartRequest request)
       return maybeResponse.get();
     }
 
-    StartResponse response =
-        StartResponse.newBuilder().setState(WorkflowExecutionState.STARTED).build();
-    context.set(WORKFLOW_EXECUTION_STATE_KEY, response);
-    return response;
+    context.set(
+        WORKFLOW_EXECUTION_STATE_KEY,
+        StartResponse.newBuilder().setState(WorkflowExecutionState.ALREADY_STARTED).build());
+    return StartResponse.newBuilder().setState(WorkflowExecutionState.STARTED).build();
   }
 
   @Override
@@ -159,6 +159,9 @@ public GetOutputResponse getOutput(KeyedContext context, OutputRequest request)
   @Override
   public void setOutput(KeyedContext context, SetOutputRequest request) throws TerminalException {
     context.set(OUTPUT_KEY, request.getOutput());
+    context.set(
+        WORKFLOW_EXECUTION_STATE_KEY,
+        StartResponse.newBuilder().setState(WorkflowExecutionState.ALREADY_COMPLETED).build());
   }
 
   @Override