Make sure we don't check the invoke time for one way calls (#381)

restatedev · Oct 16, 2024 · f0b351d · f0b351d
1 parent 4cb588d
commit f0b351d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 8 deletions.
diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/Entries.java b/sdk-core/src/main/java/dev/restate/sdk/core/Entries.java
@@ -19,6 +19,7 @@
 import io.opentelemetry.api.trace.Span;
 import java.nio.ByteBuffer;
 import java.util.Collection;
+import java.util.Objects;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -425,9 +426,10 @@ void checkEntryHeader(CallEntryMessage expected, MessageLite actual) throws Prot
       }
       CallEntryMessage actualInvoke = (CallEntryMessage) actual;
 
-      if (!(expected.getServiceName().equals(actualInvoke.getServiceName())
-          && expected.getHandlerName().equals(actualInvoke.getHandlerName())
-          && expected.getParameter().equals(actualInvoke.getParameter()))) {
+      if (!(Objects.equals(expected.getServiceName(), actualInvoke.getServiceName())
+          && Objects.equals(expected.getHandlerName(), actualInvoke.getHandlerName())
+          && Objects.equals(expected.getParameter(), actualInvoke.getParameter())
+          && Objects.equals(expected.getKey(), actualInvoke.getKey()))) {
         throw ProtocolException.entryDoesNotMatch(expected, actualInvoke);
       }
     }
@@ -452,11 +454,11 @@ public Result<R> parseCompletionResult(CompletionMessage actual) {
     }
   }
 
-  static final class BackgroundInvokeEntry extends JournalEntry<OneWayCallEntryMessage> {
+  static final class OneWayCallEntry extends JournalEntry<OneWayCallEntryMessage> {
 
-    static final BackgroundInvokeEntry INSTANCE = new BackgroundInvokeEntry();
+    static final OneWayCallEntry INSTANCE = new OneWayCallEntry();
 
-    private BackgroundInvokeEntry() {}
+    private OneWayCallEntry() {}
 
     @Override
     public void trace(OneWayCallEntryMessage expected, Span span) {
@@ -477,7 +479,17 @@ String getName(OneWayCallEntryMessage expected) {
     @Override
     void checkEntryHeader(OneWayCallEntryMessage expected, MessageLite actual)
         throws ProtocolException {
-      Util.assertEntryEquals(expected, actual);
+      if (!(actual instanceof OneWayCallEntryMessage)) {
+        throw ProtocolException.entryDoesNotMatch(expected, actual);
+      }
+      OneWayCallEntryMessage actualInvoke = (OneWayCallEntryMessage) actual;
+
+      if (!(Objects.equals(expected.getServiceName(), actualInvoke.getServiceName())
+          && Objects.equals(expected.getHandlerName(), actualInvoke.getHandlerName())
+          && Objects.equals(expected.getParameter(), actualInvoke.getParameter())
+          && Objects.equals(expected.getKey(), actualInvoke.getKey()))) {
+        throw ProtocolException.entryDoesNotMatch(expected, actualInvoke);
+      }
     }
   }
 

diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/SyscallsImpl.java b/sdk-core/src/main/java/dev/restate/sdk/core/SyscallsImpl.java
@@ -219,7 +219,7 @@ public void send(
           }
 
           this.stateMachine.processJournalEntry(
-              builder.build(), BackgroundInvokeEntry.INSTANCE, callback);
+              builder.build(), OneWayCallEntry.INSTANCE, callback);
         },
         callback);
   }