Combine onMessage and onHalfClose callbacks into single call to avoid race condition

sdk-core-impl/src/main/java/dev/restate/sdk/core/impl/ExecutorSwitchingWrappers.java

@@ -6,7 +6,6 @@
 import dev.restate.sdk.core.TypeTag;
 import dev.restate.sdk.core.syscalls.*;
 import io.grpc.MethodDescriptor;
-import io.grpc.ServerCall;
 import java.time.Duration;
 import java.util.Map;
 import java.util.concurrent.Executor;
@@ -17,8 +16,8 @@ class ExecutorSwitchingWrappers {
 
   private ExecutorSwitchingWrappers() {}
 
-  static ServerCall.Listener<MessageLite> serverCallListener(
-      ServerCall.Listener<MessageLite> sc, Executor userExecutor) {
+  static RestateServerCallListener<MessageLite> serverCallListener(
+      RestateServerCallListener<MessageLite> sc, Executor userExecutor) {
     return new ExecutorSwitchingServerCallListener(sc, userExecutor);
   }
 
@@ -27,25 +26,20 @@ static SyscallsInternal syscalls(SyscallsInternal sc, Executor syscallsExecutor)
   }
 
   private static class ExecutorSwitchingServerCallListener
-      extends ServerCall.Listener<MessageLite> {
+      implements RestateServerCallListener<MessageLite> {
 
-    private final ServerCall.Listener<MessageLite> listener;
+    private final RestateServerCallListener<MessageLite> listener;
     private final Executor userExecutor;
 
     private ExecutorSwitchingServerCallListener(
-        ServerCall.Listener<MessageLite> listener, Executor userExecutor) {
+        RestateServerCallListener<MessageLite> listener, Executor userExecutor) {
       this.listener = listener;
       this.userExecutor = userExecutor;
     }
 
     @Override
-    public void onMessage(MessageLite message) {
-      userExecutor.execute(() -> listener.onMessage(message));
-    }
-
-    @Override
-    public void onHalfClose() {
-      userExecutor.execute(listener::onHalfClose);
+    public void onMessageAndHalfClose(MessageLite message) {
+      userExecutor.execute(() -> listener.onMessageAndHalfClose(message));
     }
 
     @Override

sdk-core-impl/src/main/java/dev/restate/sdk/core/impl/ExceptionCatchingServerCallListener.java → sdk-core-impl/src/main/java/dev/restate/sdk/core/impl/GrpcServerCallListenerAdaptor.java

@@ -1,37 +1,35 @@
 package dev.restate.sdk.core.impl;
 
-import io.grpc.ForwardingServerCallListener;
 import io.grpc.Metadata;
 import io.grpc.ServerCall;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
-class ExceptionCatchingServerCallListener<ReqT, RespT>
-    extends ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT> {
+/**
+ * Adapts a {@link ServerCall.Listener} to a {@link RestateServerCallListener}.
+ *
+ * @param <ReqT> type of the request
+ * @param <RespT> type of the response
+ */
+class GrpcServerCallListenerAdaptor<ReqT, RespT> implements RestateServerCallListener<ReqT> {
 
-  private static final Logger LOG = LogManager.getLogger(ExceptionCatchingServerCallListener.class);
+  private static final Logger LOG = LogManager.getLogger(GrpcServerCallListenerAdaptor.class);
 
   private final ServerCall<ReqT, RespT> serverCall;
 
-  ExceptionCatchingServerCallListener(
+  private final ServerCall.Listener<ReqT> delegate;
+
+  GrpcServerCallListenerAdaptor(
       ServerCall.Listener<ReqT> delegate, ServerCall<ReqT, RespT> serverCall) {
-    super(delegate);
+    this.delegate = delegate;
     this.serverCall = serverCall;
   }
 
   @Override
-  public void onMessage(ReqT message) {
-    try {
-      super.onMessage(message);
-    } catch (Throwable e) {
-      closeWithException(e);
-    }
-  }
-
-  @Override
-  public void onHalfClose() {
+  public void onMessageAndHalfClose(ReqT message) {
     try {
-      super.onHalfClose();
+      delegate.onMessage(message);
+      delegate.onHalfClose();
     } catch (Throwable e) {
       closeWithException(e);
     }
@@ -40,7 +38,7 @@ public void onHalfClose() {
   @Override
   public void onCancel() {
     try {
-      super.onCancel();
+      delegate.onCancel();
     } catch (Throwable e) {
       closeWithException(e);
     }
@@ -49,7 +47,7 @@ public void onCancel() {
   @Override
   public void onComplete() {
     try {
-      super.onComplete();
+      delegate.onComplete();
     } catch (Throwable e) {
       closeWithException(e);
     }
@@ -58,7 +56,7 @@ public void onComplete() {
   @Override
   public void onReady() {
     try {
-      super.onReady();
+      delegate.onReady();
     } catch (Throwable e) {
       closeWithException(e);
     }

sdk-core-impl/src/main/java/dev/restate/sdk/core/impl/RestateGrpcServer.java

@@ -101,22 +101,23 @@ public void start() {
         stateMachine.start(
             invocationId -> {
               // Create the listener and create the decorators chain
-              ServerCall.Listener<MessageLite> listener =
+              ServerCall.Listener<MessageLite> grpcListener =
                   Contexts.interceptCall(
                       Context.current()
                           .withValue(InvocationId.INVOCATION_ID_KEY, invocationId)
                           .withValue(Syscalls.SYSCALLS_KEY, syscalls),
                       bridge,
                       new Metadata(),
                       method.getServerCallHandler());
-              listener = new ExceptionCatchingServerCallListener<>(listener, bridge);
+              RestateServerCallListener<MessageLite> restateListener =
+                  new GrpcServerCallListenerAdaptor<>(grpcListener, bridge);
               if (serverCallListenerExecutor != null) {
-                listener =
+                restateListener =
                     ExecutorSwitchingWrappers.serverCallListener(
-                        listener, serverCallListenerExecutor);
+                        restateListener, serverCallListenerExecutor);
               }
 
-              bridge.setListener(listener);
+              bridge.setListener(restateListener);
             });
       }
     };

sdk-core-impl/src/main/java/dev/restate/sdk/core/impl/RestateServerCall.java

@@ -30,7 +30,7 @@ class RestateServerCall extends ServerCall<MessageLite, MessageLite> {
   //
   // The listener reference is volatile in order to guarantee its visibility when the ownership of
   // this object is transferred through threads.
-  private volatile ServerCall.Listener<MessageLite> listener;
+  private volatile RestateServerCallListener<MessageLite> listener;
 
   // These variables don't need to be volatile as they're accessed and mutated only by
   // #setListener() and #request()
@@ -45,7 +45,7 @@ class RestateServerCall extends ServerCall<MessageLite, MessageLite> {
 
   // --- Invoked in the State machine thread
 
-  void setListener(Listener<MessageLite> listener) {
+  void setListener(RestateServerCallListener<MessageLite> listener) {
     this.listener = listener;
     this.listener.onReady();
 
@@ -158,8 +158,7 @@ private void pollInput() {
                           MessageLite message = deferredValue.toReadyResult().getResult();
 
                           LOG.trace("Read input message:\n{}", message);
-                          listener.onMessage(message);
-                          listener.onHalfClose();
+                          listener.onMessageAndHalfClose(message);
                         },
                         this::onError)),
             this::onError));

sdk-core-impl/src/main/java/dev/restate/sdk/core/impl/RestateServerCallListener.java

@@ -0,0 +1,18 @@
+package dev.restate.sdk.core.impl;
+
+/**
+ * Callbacks for incoming rpc messages.
+ *
+ * <p>This interface is strongly inspired by {@link io.grpc.ServerCall.Listener}.
+ *
+ * @param <M> type of the incoming message
+ */
+public interface RestateServerCallListener<M> {
+  void onMessageAndHalfClose(M message);
+
+  void onCancel();
+
+  void onComplete();
+
+  void onReady();
+}