From dcc009591e155ad5834d346ad1379b8b2ef2e5ce Mon Sep 17 00:00:00 2001
From: slinkydeveloper <francescoguard@gmail.com>
Date: Wed, 16 Oct 2024 16:01:38 +0200
Subject: [PATCH] Make sure we don't check the invoke time for one way calls

---
 .../java/dev/restate/sdk/core/Entries.java    | 26 ++++++++++++++-----
 .../dev/restate/sdk/core/SyscallsImpl.java    |  2 +-
 2 files changed, 20 insertions(+), 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
index b6f04dfd..010b24ce 100644
--- 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
index 1d702c3c..b657225d 100644
--- 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);
   }