chore: gracefully shutdown when error encountered

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

@@ -6,10 +6,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;
 
 /**
@@ -20,8 +22,10 @@
 
     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 sink gRPC server.
@@ -39,8 +43,10 @@
      * @param batchMapper to process the message
      */
     public Server(BatchMapper batchMapper, GRPCConfig grpcConfig) {
-        this.service = new Service(batchMapper);
+        this.shutdownSignal = new CompletableFuture<>();
+        this.service = new Service(batchMapper, this.shutdownSignal);
         this.grpcConfig = grpcConfig;
+        this.grpcServerHelper = new GrpcServerHelper();
     }
 
     /**
@@ -57,35 +63,55 @@
                 Collections.singletonMap(Constants.MAP_MODE_KEY, Constants.MAP_MODE));
 
         if (this.server == null) {
-            // create server builder
-            ServerBuilder<?> serverBuilder = GrpcServerUtils.createServerBuilder(
+            this.server = 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 socket path: " + grpcConfig.getSocketPath());
+                "server started, listening on socket path: " + 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 != null && 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 != null && server.isTerminated()) {
+                return;
+            }
+
+            if (e != null) {
+                System.err.println("*** shutting down batch map 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);
+                }
+            }
+        });
     }
 
     /**
@@ -96,7 +122,9 @@
      * @throws InterruptedException if the current thread is interrupted while waiting
      */
     public void awaitTermination() throws InterruptedException {
+        log.info("batch map server is waiting for termination");
         server.awaitTermination();
+        log.info("batch map server has terminated");
     }
 
     /**

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

@@ -35,6 +35,9 @@
     // BatchMapper instance to process the messages
     private final BatchMapper batchMapper;
 
+    // Signal to shut down the gRPC server
+    private final CompletableFuture<Void> shutdownSignal;
+
     // Applies a map function to each datum element in the stream.
     @Override
     public StreamObserver<MapOuterClass.MapRequest> mapFn(StreamObserver<MapOuterClass.MapResponse> responseObserver) {
@@ -93,8 +96,9 @@
                         datumStream.writeMessage(constructHandlerDatum(mapRequest));
                     }
                 } catch (Exception e) {
-                    log.error("Encountered an error in batch map", e);
-                    responseObserver.onError(Status.UNKNOWN
+                    log.error("Encountered an error in batch map onNext - {}", e.getMessage());
+                    shutdownSignal.completeExceptionally(e);
+                    responseObserver.onError(Status.INTERNAL
                             .withDescription(e.getMessage())
                             .withCause(e)
                             .asException());
@@ -104,11 +108,12 @@
             // Called when an error occurs
             @Override
             public void onError(Throwable throwable) {
-                log.error("Error Encountered in batchMap Stream", throwable);
-                var status = Status.UNKNOWN
+                log.error("Error Encountered in batchMap Stream - {}", throwable.getMessage());
+                shutdownSignal.completeExceptionally(throwable);
+                responseObserver.onError(Status.INTERNAL
                         .withDescription(throwable.getMessage())
-                        .withCause(throwable);
-                responseObserver.onError(status.asException());
+                        .withCause(throwable)
+                        .asException());
             }
 
             // Called when the client has finished sending requests

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,12 @@
     }
 
     private void handleFailure(Exception e) {
-        responseObserver.onError(Status.UNKNOWN
+        log.error("Encountered error in mapFn - {}", e.getMessage());
+        shutdownSignal.completeExceptionally(e);
+        responseObserver.onError(Status.INTERNAL
                 .withDescription(e.getMessage())
                 .withCause(e)
                 .asException());
-        failureFuture.completeExceptionally(e);
     }
 
     private void sendResponse(MapOuterClass.MapResponse mapResponse) {
@@ -118,8 +119,8 @@
     // 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 +130,8 @@
         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 @@
 
     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 @@
      * @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 @@
         }
 
         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 != null && 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 @@
      * @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;