map streamer

Signed-off-by: Keran Yang <yangkr920208@gmail.com>

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

@@ -80,7 +80,7 @@ public void start() throws Exception {
         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()) {
+            if (server != null && server.isTerminated()) {
                 return;
             }
             try {
@@ -95,7 +95,7 @@ public void start() throws Exception {
 
         // if there are any exceptions, shutdown the server gracefully.
         shutdownSignal.whenCompleteAsync((v, e) -> {
-            if (server.isTerminated()) {
+            if (server != null && server.isTerminated()) {
                 return;
             }
 

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

@@ -86,7 +86,7 @@ public void start() throws Exception {
         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()) {
+            if (server != null && server.isTerminated()) {
                 return;
             }
             try {

src/main/java/io/numaproj/numaflow/mapstreamer/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 @@ 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 sink gRPC server.
@@ -39,8 +43,10 @@ public Server(MapStreamer mapStreamer) {
      * @param mapStreamer to process the message
      */
     public Server(MapStreamer mapStreamer, GRPCConfig grpcConfig) {
-        this.service = new Service(mapStreamer);
+        this.shutdownSignal = new CompletableFuture<>();
+        this.service = new Service(mapStreamer, this.shutdownSignal);
         this.grpcConfig = grpcConfig;
+        this.grpcServerHelper = new GrpcServerHelper();
     }
 
     /**
@@ -57,35 +63,55 @@ public void start() throws Exception {
                 Collections.singletonMap(Constants.MAP_MODE_KEY, Constants.MAP_MODE));
 
         if (this.server == null) {
-            // create server builder
-            ServerBuilder<?> 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 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");
+            System.err.println("*** shutting down map streamer 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 map streamer 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 @@ public void start() throws Exception {
      * @throws InterruptedException if the current thread is interrupted while waiting
      */
     public void awaitTermination() throws InterruptedException {
+        log.info("sink server is waiting for termination");
         server.awaitTermination();
+        log.info("sink server has terminated");
     }
 
     /**

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

@@ -9,6 +9,7 @@
 import lombok.extern.slf4j.Slf4j;
 
 import java.time.Instant;
+import java.util.concurrent.CompletableFuture;
 
 import static io.numaproj.numaflow.map.v1.MapGrpc.getMapFnMethod;
 
@@ -17,6 +18,7 @@
 class Service extends MapGrpc.MapImplBase {
 
     private final MapStreamer mapStreamer;
+    private final CompletableFuture<Void> shutdownSignal;
 
     @Override
     public StreamObserver<MapOuterClass.MapRequest> mapFn(StreamObserver<MapOuterClass.MapResponse> responseObserver) {
@@ -57,9 +59,11 @@ public void onNext(MapOuterClass.MapRequest request) {
                             constructHandlerDatum(request),
                             new OutputObserverImpl(responseObserver));
                 } catch (Exception e) {
-                    log.error("Error processing message", e);
-                    responseObserver.onError(Status.UNKNOWN
+                    log.error("Encountered error in mapFn onNext - {}", e.getMessage());
+                    shutdownSignal.completeExceptionally(e);
+                    responseObserver.onError(Status.INTERNAL
                             .withDescription(e.getMessage())
+                            .withCause(e)
                             .asException());
                     return;
                 }
@@ -76,11 +80,12 @@ public void onNext(MapOuterClass.MapRequest request) {
 
             @Override
             public void onError(Throwable throwable) {
-                log.error("Error Encountered in mapStream Stream", throwable);
-                var status = Status.UNKNOWN
+                log.error("Encountered error in mapStream Stream - {}", throwable.getMessage());
+                shutdownSignal.completeExceptionally(throwable);
+                responseObserver.onError(Status.UNKNOWN
                         .withDescription(throwable.getMessage())
-                        .withCause(throwable);
-                responseObserver.onError(status.asException());
+                        .withCause(throwable)
+                        .asException());
             }
 
             @Override

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

@@ -79,7 +79,7 @@ public void start() throws Exception {
         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()) {
+            if (server != null && server.isTerminated()) {
                 return;
             }
             try {

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

@@ -81,7 +81,7 @@ public void start() throws Exception {
         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 sink gRPC server since JVM is shutting down");
-            if (server.isTerminated()) {
+            if (server != null && server.isTerminated()) {
                 return;
             }
             try {

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

@@ -128,7 +128,7 @@ public void TestMapStreamerErr() {
             fail("Should have thrown an exception");
         } catch (Exception e) {
             assertEquals(
-                    "io.grpc.StatusRuntimeException: UNKNOWN: unknown exception",
+                    "io.grpc.StatusRuntimeException: INTERNAL: unknown exception",
                     e.getMessage());
         }
     }