Small things

Author: slinkydeveloper

sdk-api-kotlin/src/main/kotlin/dev/restate/sdk/kotlin/ContextImpl.kt

@@ -123,8 +123,8 @@ internal constructor(
       retryPolicy: RetryPolicy?,
       block: suspend () -> T
   ): DurableFuture<T> {
-    var serde: Serde<T> = resolveSerde(typeTag)
-    var coroutineCtx = currentCoroutineContext()
+    val serde: Serde<T> = resolveSerde(typeTag)
+    val coroutineCtx = currentCoroutineContext()
     val javaRetryPolicy =
         retryPolicy?.let {
           dev.restate.sdk.common.RetryPolicy.exponential(

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

@@ -173,6 +173,12 @@ public RequestProcessor processorForRequest(
 
     String fullyQualifiedServiceMethod = serviceName + "/" + handlerName;
 
+    // If we got it, set it
+    String invocationIdHeader = headersAccessor.get("x-restate-invocation-id");
+    if (invocationIdHeader != null) {
+      loggingContextSetter.set(LoggingContextSetter.INVOCATION_ID_KEY, invocationIdHeader);
+    }
+
     // Instantiate state machine
     StateMachine stateMachine = StateMachine.init(headersAccessor, loggingContextSetter);
 

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

@@ -42,7 +42,6 @@ class HandlerContextImpl implements HandlerContextInternal {
   private final @Nullable String objectKey;
   private final String fullyQualifiedHandlerName;
 
-  private CompletableFuture<Void> nextProcessedRun;
   private final List<AsyncResultInternal<String>> invocationIdsToCancel;
   private final HashMap<Integer, Consumer<RunCompleter>> scheduledRuns;
 
@@ -416,7 +415,6 @@ private void pollAsyncResultInner(AsyncResultInternal<?> asyncResult) {
         // Let it loop now
       } else if (response instanceof StateMachine.DoProgressResponse.ReadFromInput
           || response instanceof StateMachine.DoProgressResponse.WaitingPendingRun) {
-//        LOG.info("Gonna need to wait here {}", response);
         this.stateMachine.onNextEvent(
             () -> {
               LOG.info("Triggered after wait {}", response);

sdk-core/src/main/java/dev/restate/sdk/core/statemachine/StateMachineImpl.java

@@ -74,16 +74,18 @@ public CompletableFuture<Void> waitForReady() {
 
   @Override
   public void onNextEvent(Runnable runnable) {
-//    LOG.info("Registering next event signal");
     this.nextEventListener =
         () -> {
           this.nextEventListener.run();
           runnable.run();
         };
+    // Trigger this now
+    if (this.stateContext.isInputClosed()) {
+      this.triggerNextEventSignal();
+    }
   }
 
   private void triggerNextEventSignal() {
-//    LOG.info("Triggering next event signal");
     Runnable listener = this.nextEventListener;
     this.nextEventListener = () -> {};
     listener.run();
@@ -163,6 +165,7 @@ public void onComplete() {
       this.stateContext.getCurrentState().onInputClosed(this.stateContext);
     } catch (Throwable e) {
       this.onError(e);
+      return;
     }
     this.triggerNextEventSignal();
     this.cancelInputSubscription();
@@ -550,9 +553,11 @@ public void proposeRunCompletion(int handle, Slice value) {
     LOG.debug("Executing 'Run completed with success'");
     try {
       this.stateContext.getCurrentState().proposeRunCompletion(handle, value, this.stateContext);
-    } finally {
-      this.triggerNextEventSignal();
+    } catch (Throwable e) {
+      this.onError(e);
+      return;
     }
+    this.triggerNextEventSignal();
   }
 
   @Override
@@ -566,9 +571,11 @@ public void proposeRunCompletion(
       this.stateContext
           .getCurrentState()
           .proposeRunCompletion(handle, exception, attemptDuration, retryPolicy, this.stateContext);
-    } finally {
-      this.triggerNextEventSignal();
+    } catch (Throwable e) {
+      this.onError(e);
+      return;
     }
+    this.triggerNextEventSignal();
   }
 
   @Override