Fix side effect guard (#83)

restatedev · Apr 24, 2023 · ca4617d · ca4617d
1 parent 389f77f
commit ca4617d
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/sdk-core-impl/src/main/java/dev/restate/sdk/core/impl/InvocationStateMachine.java b/sdk-core-impl/src/main/java/dev/restate/sdk/core/impl/InvocationStateMachine.java
@@ -305,7 +305,6 @@ void enterSideEffectJournalEntry(
       Runnable noEntryCallback,
       Consumer<Throwable> failureCallback) {
     checkInsideSideEffectGuard();
-    this.insideSideEffect = true;
     if (this.state == State.CLOSED) {
       failureCallback.accept(SuspendedException.INSTANCE);
     } else if (this.state == State.REPLAYING) {
@@ -322,6 +321,7 @@ void enterSideEffectJournalEntry(
           new SideEffectAckPublisher.OnEnterSideEffectCallback() {
             @Override
             public void onEnter() {
+              insideSideEffect = true;
               if (span.isRecording()) {
                 traceFn.accept(span);
               }

diff --git a/sdk-core-impl/src/test/java/dev/restate/sdk/core/impl/SideEffectTest.java b/sdk-core-impl/src/test/java/dev/restate/sdk/core/impl/SideEffectTest.java
@@ -98,6 +98,24 @@ public void greet(GreetingRequest request, StreamObserver<GreetingResponse> resp
     }
   }
 
+  private static class SideEffectThenAwakeable extends GreeterGrpc.GreeterImplBase
+      implements RestateBlockingService {
+
+    @Override
+    public void greet(GreetingRequest request, StreamObserver<GreetingResponse> responseObserver) {
+      RestateContext ctx = restateContext();
+
+      ctx.sideEffect(
+          () -> {
+            throw new IllegalStateException("This should be replayed");
+          });
+      ctx.awakeable(TypeTag.BYTES).await();
+
+      responseObserver.onNext(GreetingResponse.newBuilder().setMessage("Hello").build());
+      responseObserver.onCompleted();
+    }
+  }
+
   @Override
   Stream<TestDefinition> definitions() {
     return Stream.of(
@@ -153,6 +171,13 @@ Stream<TestDefinition> definitions() {
         testInvocation(new SideEffectGuard(), GreeterGrpc.getGreetMethod())
             .withInput(startMessage(1), inputMessage(GreetingRequest.newBuilder().setName("Till")))
             .usingAllThreadingModels()
-            .assertingFailure(ProtocolException.class));
+            .assertingFailure(ProtocolException.class),
+        testInvocation(new SideEffectThenAwakeable(), GreeterGrpc.getGreetMethod())
+            .withInput(
+                startMessage(2),
+                inputMessage(GreetingRequest.newBuilder().setName("Till")),
+                Java.SideEffectEntryMessage.newBuilder().setValue(ByteString.copyFromUtf8("")))
+            .usingAllThreadingModels()
+            .expectingOutput(awakeable(), suspensionMessage(2)));
   }
 }