map

Signed-off-by: Keran Yang <[email protected]>

src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java

@@ -46,29 +46,29 @@
 class MapSupervisorActor extends AbstractActor {
     private final Mapper mapper;
     private final StreamObserver<MapOuterClass.MapResponse> responseObserver;
-    private final CompletableFuture<Void> failureFuture;
+    private final CompletableFuture<Void> shutdownSignal;
 
     public MapSupervisorActor(
             Mapper mapper,
             StreamObserver<MapOuterClass.MapResponse> responseObserver,
             CompletableFuture<Void> failureFuture) {
         this.mapper = mapper;
         this.responseObserver = responseObserver;
-        this.failureFuture = failureFuture;
+        this.shutdownSignal = failureFuture;
     }
 
     public static Props props(
             Mapper mapper,
             StreamObserver<MapOuterClass.MapResponse> responseObserver,
-            CompletableFuture<Void> failureFuture) {
-        return Props.create(MapSupervisorActor.class, mapper, responseObserver, failureFuture);
+            CompletableFuture<Void> shutdownSignal) {
+        return Props.create(MapSupervisorActor.class, mapper, responseObserver, shutdownSignal);
     }
 
     @Override
     public void preRestart(Throwable reason, Optional<Object> message) {
         log.debug("supervisor pre restart was executed");
-        failureFuture.completeExceptionally(reason);
-        responseObserver.onError(Status.UNKNOWN
+        shutdownSignal.completeExceptionally(reason);
+        responseObserver.onError(Status.INTERNAL
                 .withDescription(reason.getMessage())
                 .withCause(reason)
                 .asException());
@@ -93,11 +93,11 @@ public Receive createReceive() {
     }
 
     private void handleFailure(Exception e) {
-        responseObserver.onError(Status.UNKNOWN
+        shutdownSignal.completeExceptionally(e);
+        responseObserver.onError(Status.INTERNAL
                 .withDescription(e.getMessage())
                 .withCause(e)
                 .asException());
-        failureFuture.completeExceptionally(e);
     }
 
     private void sendResponse(MapOuterClass.MapResponse mapResponse) {
@@ -118,8 +118,8 @@ private void processRequest(MapOuterClass.MapRequest mapRequest) {
     // to make sure no messages are lost
     private void handleDeadLetters(AllDeadLetters deadLetter) {
         log.debug("got a dead letter, stopping the execution");
-        responseObserver.onError(Status.UNKNOWN.withDescription("dead letters").asException());
-        failureFuture.completeExceptionally(new Throwable("dead letters"));
+        shutdownSignal.completeExceptionally(new Throwable("dead letters"));
+        responseObserver.onError(Status.INTERNAL.withDescription("dead letters").asException());
         getContext().getSystem().stop(getSelf());
     }
 
@@ -129,8 +129,8 @@ public SupervisorStrategy supervisorStrategy() {
         return new AllForOneStrategy(
                 DeciderBuilder
                         .match(Exception.class, e -> {
-                            failureFuture.completeExceptionally(e);
-                            responseObserver.onError(Status.UNKNOWN
+                            shutdownSignal.completeExceptionally(e);
+                            responseObserver.onError(Status.INTERNAL
                                     .withDescription(e.getMessage())
                                     .withCause(e)
                                     .asException());

src/main/java/io/numaproj/numaflow/mapper/Server.java

@@ -5,10 +5,12 @@
 import io.numaproj.numaflow.info.ContainerType;
 import io.numaproj.numaflow.info.ServerInfoAccessor;
 import io.numaproj.numaflow.info.ServerInfoAccessorImpl;
+import io.numaproj.numaflow.shared.GrpcServerHelper;
 import io.numaproj.numaflow.shared.GrpcServerUtils;
 import lombok.extern.slf4j.Slf4j;
 
 import java.util.Collections;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -19,8 +21,10 @@ public class Server {
 
     private final GRPCConfig grpcConfig;
     private final Service service;
+    private final CompletableFuture<Void> shutdownSignal;
     private final ServerInfoAccessor serverInfoAccessor = new ServerInfoAccessorImpl(new ObjectMapper());
     private io.grpc.Server server;
+    private final GrpcServerHelper grpcServerHelper;
 
     /**
      * constructor to create gRPC server.
@@ -38,8 +42,10 @@ public Server(Mapper mapper) {
      * @param mapper to process the message
      */
     public Server(Mapper mapper, GRPCConfig grpcConfig) {
-        this.service = new Service(mapper);
+        this.shutdownSignal = new CompletableFuture<>();
+        this.service = new Service(mapper, this.shutdownSignal);
         this.grpcConfig = grpcConfig;
+        this.grpcServerHelper = new GrpcServerHelper();
     }
 
     /**
@@ -61,39 +67,57 @@ public void start() throws Exception {
         }
 
         if (this.server == null) {
-            ServerBuilder<?> serverBuilder;
-            // create server builder for domain socket server
-            serverBuilder = GrpcServerUtils.createServerBuilder(
+            this.server = this.grpcServerHelper.createServer(
                     grpcConfig.getSocketPath(),
                     grpcConfig.getMaxMessageSize(),
                     grpcConfig.isLocal(),
-                    grpcConfig.getPort());
-
-            // build server
-            this.server = serverBuilder
-                    .addService(this.service)
-                    .build();
+                    grpcConfig.getPort(),
+                    this.service);
         }
 
-        // start server
         server.start();
 
         log.info(
-                "Server started, listening on {}",
+                "server started, listening on {}",
                 grpcConfig.isLocal() ?
                         "localhost:" + grpcConfig.getPort() : grpcConfig.getSocketPath());
 
-        // register shutdown hook
+        // register shutdown hook to gracefully shut down the server
         Runtime.getRuntime().addShutdownHook(new Thread(() -> {
             // Use stderr here since the logger may have been reset by its JVM shutdown hook.
             System.err.println("*** shutting down gRPC server since JVM is shutting down");
+            if (server.isTerminated()) {
+                return;
+            }
             try {
                 Server.this.stop();
+                log.info("gracefully shutting down event loop groups");
+                this.grpcServerHelper.gracefullyShutdownEventLoopGroups();
             } catch (InterruptedException e) {
                 Thread.interrupted();
                 e.printStackTrace(System.err);
             }
         }));
+
+        // if there are any exceptions, shutdown the server gracefully.
+        shutdownSignal.whenCompleteAsync((v, e) -> {
+            if (server.isTerminated()) {
+                return;
+            }
+
+            if (e != null) {
+                System.err.println("*** shutting down mapper gRPC server because of an exception - " + e.getMessage());
+                try {
+                    log.info("stopping server");
+                    Server.this.stop();
+                    log.info("gracefully shutting down event loop groups");
+                    this.grpcServerHelper.gracefullyShutdownEventLoopGroups();
+                } catch (InterruptedException ex) {
+                    Thread.interrupted();
+                    ex.printStackTrace(System.err);
+                }
+            }
+        });
     }
 
     /**
@@ -104,7 +128,9 @@ public void start() throws Exception {
      * @throws InterruptedException if the current thread is interrupted while waiting
      */
     public void awaitTermination() throws InterruptedException {
+        log.info("mapper server is waiting for termination");
         server.awaitTermination();
+        log.info("mapper server has terminated");
     }
 
     /**

src/main/java/io/numaproj/numaflow/mapper/Service.java

@@ -21,18 +21,7 @@ class Service extends MapGrpc.MapImplBase {
     public static final ActorSystem mapperActorSystem = ActorSystem.create("mapper");
 
     private final Mapper mapper;
-
-    // TODO we need to propagate the exception all the way up and shutdown the server.
-    static void handleFailure(
-            CompletableFuture<Void> failureFuture) {
-        new Thread(() -> {
-            try {
-                failureFuture.get();
-            } catch (Exception e) {
-                e.printStackTrace();
-            }
-        }).start();
-    }
+    private final CompletableFuture<Void> shutdownSignal;
 
     @Override
     public StreamObserver<MapOuterClass.MapRequest> mapFn(final StreamObserver<MapOuterClass.MapResponse> responseObserver) {
@@ -43,13 +32,9 @@ public StreamObserver<MapOuterClass.MapRequest> mapFn(final StreamObserver<MapOu
                     responseObserver);
         }
 
-        CompletableFuture<Void> failureFuture = new CompletableFuture<>();
-
-        handleFailure(failureFuture);
-
         // create a MapSupervisorActor that processes the map requests.
         ActorRef mapSupervisorActor = mapperActorSystem
-                .actorOf(MapSupervisorActor.props(mapper, responseObserver, failureFuture));
+                .actorOf(MapSupervisorActor.props(mapper, responseObserver, shutdownSignal));
 
         return new StreamObserver<>() {
             private boolean handshakeDone = false;

src/test/java/io/numaproj/numaflow/mapper/ServerErrTest.java

@@ -129,7 +129,7 @@ public void testMapperFailure() {
             fail("Expected exception not thrown");
         } catch (Exception e) {
             assertEquals(
-                    "io.grpc.StatusRuntimeException: UNKNOWN: unknown exception",
+                    "io.grpc.StatusRuntimeException: INTERNAL: unknown exception",
                     e.getMessage());
         }
     }