Use the StartMessage.partial_state flag, Remove partial_state flag pa…

…rsing
restatedev · Sep 6, 2023 · 0349058 · 0349058
1 parent 2a01b5e
commit 0349058
Show file tree

Hide file tree

Showing 7 changed files with 24 additions and 32 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
@@ -111,11 +111,10 @@ public void onSubscribe(Flow.Subscription subscription) {
 
   @Override
   public void onNext(InvocationFlow.InvocationInput invocationInput) {
-    MessageHeader header = invocationInput.header();
     MessageLite msg = invocationInput.message();
     LOG.trace("Received input message {} {}", msg.getClass(), msg);
     if (this.state == State.WAITING_START) {
-      this.onStart(header, msg);
+      this.onStart(msg);
     } else if (msg instanceof Protocol.CompletionMessage) {
       // We check the instance rather than the state, because the user code might still be
       // replaying, but the network layer is already past it and is receiving completions from the
@@ -154,7 +153,7 @@ void start(Consumer<InvocationId> afterStartCallback) {
     this.inputSubscription.request(1);
   }
 
-  void onStart(MessageHeader header, MessageLite msg) {
+  void onStart(MessageLite msg) {
     if (!(msg instanceof Protocol.StartMessage)) {
       this.fail(ProtocolException.unexpectedMessage(Protocol.StartMessage.class, msg));
       return;
@@ -169,7 +168,7 @@ void onStart(MessageHeader header, MessageLite msg) {
     // Set up the state cache
     this.localStateStorage =
         new LocalStateStorage(
-            header.hasFlag(MessageHeader.PARTIAL_STATE_FLAG),
+            startMessage.getPartialState(),
             startMessage.getStateMapList().stream()
                 .collect(
                     Collectors.toMap(

diff --git a/sdk-core-impl/src/main/java/dev/restate/sdk/core/impl/MessageHeader.java b/sdk-core-impl/src/main/java/dev/restate/sdk/core/impl/MessageHeader.java
@@ -6,7 +6,6 @@
 
 public class MessageHeader {
 
-  public static final short PARTIAL_STATE_FLAG = 0x0400;
   private static final short DONE_FLAG = 0x0001;
   private static final short REQUIRES_ACK_FLAG = 0x0001;
 
@@ -36,10 +35,6 @@ public long encode() {
     return res;
   }
 
-  public boolean hasFlag(short flag) {
-    return (this.flags & flag) > 0;
-  }
-
   public MessageHeader copyWithFlags(short flag) {
     return new MessageHeader(type, flag, length);
   }

diff --git a/sdk-core-impl/src/test/java/dev/restate/sdk/core/impl/EagerStateTest.java b/sdk-core-impl/src/test/java/dev/restate/sdk/core/impl/EagerStateTest.java
@@ -77,8 +77,6 @@ public void greet(GreetingRequest request, StreamObserver<GreetingResponse> resp
     }
   }
 
-  private static final short PARTIAL_STATE = PARTIAL_STATE_FLAG;
-  private static final short COMPLETE_STATE = 0;
   private static final Map.Entry<String, String> STATE_FRANCESCO = entry("STATE", "Francesco");
   private static final MessageLite INPUT_TILL = inputMessage(greetingRequest("Till"));
   private static final MessageLite GET_STATE_FRANCESCO = getStateMessage("STATE", "Francesco");
@@ -94,38 +92,33 @@ public void greet(GreetingRequest request, StreamObserver<GreetingResponse> resp
   Stream<TestDefinition> definitions() {
     return Stream.of(
         testInvocation(new GetEmpty(), GreeterGrpc.getGreetMethod())
-            .withInput(COMPLETE_STATE, startMessage(1))
-            .withInput(INPUT_TILL)
+            .withInput(startMessage(1).setPartialState(false), INPUT_TILL)
             .usingAllThreadingModels()
             .expectingOutput(getStateEmptyMessage("STATE"), outputMessage(greetingResponse("true")))
             .named("With complete state"),
         testInvocation(new GetEmpty(), GreeterGrpc.getGreetMethod())
-            .withInput(PARTIAL_STATE, startMessage(1))
-            .withInput(INPUT_TILL)
+            .withInput(startMessage(1).setPartialState(true), INPUT_TILL)
             .usingAllThreadingModels()
             .expectingOutput(getStateMessage("STATE"), suspensionMessage(1))
             .named("With partial state"),
         testInvocation(new GetEmpty(), GreeterGrpc.getGreetMethod())
-            .withInput(PARTIAL_STATE, startMessage(2))
-            .withInput(INPUT_TILL, getStateEmptyMessage("STATE"))
+            .withInput(
+                startMessage(2).setPartialState(true), INPUT_TILL, getStateEmptyMessage("STATE"))
             .usingAllThreadingModels()
             .expectingOutput(outputMessage(greetingResponse("true")))
             .named("Resume with partial state"),
         testInvocation(new Get(), GreeterGrpc.getGreetMethod())
-            .withInput(COMPLETE_STATE, startMessage(1, STATE_FRANCESCO))
-            .withInput(INPUT_TILL)
+            .withInput(startMessage(1, STATE_FRANCESCO).setPartialState(false), INPUT_TILL)
             .usingAllThreadingModels()
             .expectingOutput(GET_STATE_FRANCESCO, OUTPUT_FRANCESCO)
             .named("With complete state"),
         testInvocation(new Get(), GreeterGrpc.getGreetMethod())
-            .withInput(PARTIAL_STATE, startMessage(1, STATE_FRANCESCO))
-            .withInput(INPUT_TILL)
+            .withInput(startMessage(1, STATE_FRANCESCO).setPartialState(true), INPUT_TILL)
             .usingAllThreadingModels()
             .expectingOutput(GET_STATE_FRANCESCO, OUTPUT_FRANCESCO)
             .named("With partial state"),
         testInvocation(new Get(), GreeterGrpc.getGreetMethod())
-            .withInput(PARTIAL_STATE, startMessage(1))
-            .withInput(INPUT_TILL)
+            .withInput(startMessage(1).setPartialState(true), INPUT_TILL)
             .usingAllThreadingModels()
             .expectingOutput(getStateMessage("STATE"), suspensionMessage(1))
             .named("With partial state without the state entry"),
@@ -139,8 +132,10 @@ Stream<TestDefinition> definitions() {
                 OUTPUT_FRANCESCO_TILL)
             .named("With state in the state_map"),
         testInvocation(new GetAppendAndGet(), GreeterGrpc.getGreetMethod())
-            .withInput(PARTIAL_STATE, startMessage(1))
-            .withInput(INPUT_TILL, completionMessage(1, "Francesco"))
+            .withInput(
+                startMessage(1).setPartialState(true),
+                INPUT_TILL,
+                completionMessage(1, "Francesco"))
             .usingAllThreadingModels()
             .expectingOutput(
                 getStateMessage("STATE"),
@@ -158,8 +153,10 @@ Stream<TestDefinition> definitions() {
                 OUTPUT_FRANCESCO)
             .named("With state in the state_map"),
         testInvocation(new GetClearAndGet(), GreeterGrpc.getGreetMethod())
-            .withInput(PARTIAL_STATE, startMessage(1))
-            .withInput(INPUT_TILL, completionMessage(1, "Francesco"))
+            .withInput(
+                startMessage(1).setPartialState(true),
+                INPUT_TILL,
+                completionMessage(1, "Francesco"))
             .usingAllThreadingModels()
             .expectingOutput(
                 getStateMessage("STATE"),

diff --git a/sdk-core-impl/src/test/java/dev/restate/sdk/core/impl/ProtoUtils.java b/sdk-core-impl/src/test/java/dev/restate/sdk/core/impl/ProtoUtils.java
@@ -26,8 +26,7 @@ public class ProtoUtils {
    */
   public static MessageHeader headerFromMessage(MessageLite msg) {
     if (msg instanceof Protocol.StartMessage) {
-      return new MessageHeader(
-          MessageType.StartMessage, MessageHeader.PARTIAL_STATE_FLAG, msg.getSerializedSize());
+      return new MessageHeader(MessageType.StartMessage, (short) 0, msg.getSerializedSize());
     } else if (msg instanceof Protocol.CompletionMessage) {
       return new MessageHeader(MessageType.CompletionMessage, (short) 0, msg.getSerializedSize());
     }
@@ -38,7 +37,8 @@ public static Protocol.StartMessage.Builder startMessage(int entries) {
     return Protocol.StartMessage.newBuilder()
         .setId(ByteString.copyFromUtf8("abc"))
         .setDebugId("abc")
-        .setKnownEntries(entries);
+        .setKnownEntries(entries)
+        .setPartialState(true);
   }
 
   @SafeVarargs

diff --git a/sdk-lambda/src/test/java/dev/restate/sdk/lambda/LambdaHandlerTest.java b/sdk-lambda/src/test/java/dev/restate/sdk/lambda/LambdaHandlerTest.java
@@ -48,6 +48,7 @@ public void testInvoke(String serviceName) throws IOException {
                         .setDebugId("123")
                         .setId(ByteString.copyFromUtf8("123"))
                         .setKnownEntries(1)
+                        .setPartialState(true)
                         .build(),
                     Protocol.PollInputStreamEntryMessage.newBuilder()
                         .setValue(

diff --git a/sdk-testing/src/main/java/dev/restate/sdk/testing/PublishSubscription.java b/sdk-testing/src/main/java/dev/restate/sdk/testing/PublishSubscription.java
@@ -33,8 +33,7 @@ public void cancel() {}
 
   static MessageHeader headerFromMessage(MessageLite msg) {
     if (msg instanceof Protocol.StartMessage) {
-      return new MessageHeader(
-          MessageType.StartMessage, MessageHeader.PARTIAL_STATE_FLAG, msg.getSerializedSize());
+      return new MessageHeader(MessageType.StartMessage, (short) 0, msg.getSerializedSize());
     } else if (msg instanceof Protocol.CompletionMessage) {
       return new MessageHeader(MessageType.CompletionMessage, (short) 0, msg.getSerializedSize());
     }

diff --git a/sdk-testing/src/main/java/dev/restate/sdk/testing/TestRestateRuntime.java b/sdk-testing/src/main/java/dev/restate/sdk/testing/TestRestateRuntime.java
@@ -109,6 +109,7 @@ private void handle(
             .setDebugId(invocationId)
             .setId(ByteString.copyFromUtf8(invocationId))
             .setKnownEntries(1)
+            .setPartialState(true)
             .build();
     List<MessageLite> inputMessages = List.of(startMessage, inputMessage);