chore: gracefully shutdown when error encountered

src/main/java/io/numaproj/numaflow/shared/GrpcServerHelper.java

@@ -0,0 +1,110 @@
+package io.numaproj.numaflow.shared;
+
+import io.grpc.BindableService;
+import io.grpc.Context;
+import io.grpc.Contexts;
+import io.grpc.ForwardingServerCallListener;
+import io.grpc.Server;
+import io.grpc.ServerBuilder;
+import io.grpc.ServerCall;
+import io.grpc.ServerCallHandler;
+import io.grpc.ServerInterceptor;
+import io.grpc.Status;
+import io.grpc.netty.NettyServerBuilder;
+import io.netty.channel.EventLoopGroup;
+import io.netty.channel.unix.DomainSocketAddress;
+
+import static io.numaproj.numaflow.shared.GrpcServerUtils.DATUM_METADATA_WIN_END;
+import static io.numaproj.numaflow.shared.GrpcServerUtils.DATUM_METADATA_WIN_START;
+import static io.numaproj.numaflow.shared.GrpcServerUtils.WINDOW_END_TIME;
+import static io.numaproj.numaflow.shared.GrpcServerUtils.WINDOW_START_TIME;
+
+public class GrpcServerHelper {
+    private EventLoopGroup bossEventLoopGroup;
+    private EventLoopGroup workerEventLoopGroup;
+
+    public void gracefullyShutdownEventLoopGroups() {
+        if (this.bossEventLoopGroup != null) {
+            this.bossEventLoopGroup.shutdownGracefully();
+        }
+        if (this.workerEventLoopGroup != null) {
+            this.workerEventLoopGroup.shutdownGracefully();
+        }
+    }
+
+    public Server createServer(
+            String socketPath,
+            int maxMessageSize,
+            boolean isLocal,
+            int port,
+            BindableService service) {
+        ServerInterceptor interceptor = new ServerInterceptor() {
+            @Override
+            public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
+                    ServerCall<ReqT, RespT> call,
+                    io.grpc.Metadata headers,
+                    ServerCallHandler<ReqT, RespT> next) {
+
+                final var context =
+                        Context.current().withValues(
+                                WINDOW_START_TIME,
+                                headers.get(DATUM_METADATA_WIN_START),
+                                WINDOW_END_TIME,
+                                headers.get(DATUM_METADATA_WIN_END));
+                ServerCall.Listener<ReqT> listener = Contexts.interceptCall(
+                        context,
+                        call,
+                        headers,
+                        next);
+                return new ForwardingServerCallListener.SimpleForwardingServerCallListener<>(
+                        listener) {
+                    @Override
+                    public void onHalfClose() {
+                        try {
+                            super.onHalfClose();
+                        } catch (RuntimeException ex) {
+                            handleException(ex, call, headers);
+                            throw ex;
+                        }
+                    }
+
+                    private void handleException(
+                            RuntimeException e,
+                            ServerCall<ReqT, RespT> serverCall,
+                            io.grpc.Metadata headers) {
+                        // Currently, we only have application level exceptions.
+                        // Translate it to UNKNOWN status.
+                        var status = Status.UNKNOWN.withDescription(e.getMessage()).withCause(e);
+                        var newStatus = Status.fromThrowable(status.asException());
+                        serverCall.close(newStatus, headers);
+                        e.printStackTrace();
+                        System.exit(1);
+                    }
+                };
+            }
+        };
+
+        if (isLocal) {
+            return ServerBuilder.forPort(port)
+                    .maxInboundMessageSize(maxMessageSize)
+                    .intercept(interceptor)
+                    .addService(service)
+                    .build();
+        }
+
+        this.bossEventLoopGroup = GrpcServerUtils.createEventLoopGroup(1, "netty-boss");
+        this.workerEventLoopGroup = GrpcServerUtils.createEventLoopGroup(
+                ThreadUtils.INSTANCE.availableProcessors(),
+                "netty-worker");
+
+        return NettyServerBuilder
+                .forAddress(new DomainSocketAddress(socketPath))
+                .channelType(GrpcServerUtils.getChannelTypeClass())
+                .maxInboundMessageSize(maxMessageSize)
+                .bossEventLoopGroup(this.bossEventLoopGroup)
+                .workerEventLoopGroup(this.workerEventLoopGroup)
+                .intercept(interceptor)
+                .addService(service)
+                .build();
+    }
+}

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

@@ -5,9 +5,11 @@
 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.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -18,8 +20,10 @@
 
     private final GRPCConfig grpcConfig;
     private final Service service;
+    public 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.
@@ -37,8 +41,10 @@
      * @param sinker sink to process the message
      */
     public Server(Sinker sinker, GRPCConfig grpcConfig) {
-        this.service = new Service(sinker);
+        this.shutdownSignal = new CompletableFuture<>();
+        this.service = new Service(sinker, this.shutdownSignal);
         this.grpcConfig = grpcConfig;
+        this.grpcServerHelper = new GrpcServerHelper();
     }
 
     /**
@@ -56,38 +62,57 @@
         }
 
         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 {}",
+                "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");
+            System.err.println("*** shutting down sink 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 sink 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);
+                }
+            }
+        });
     }
 
     /**
@@ -98,7 +123,9 @@
      * @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 is terminated");
     }
 
     /**
@@ -108,14 +135,17 @@
      * @throws InterruptedException if shutdown is interrupted
      */
     public void stop() throws InterruptedException {
+        log.info("server.stop started. Shutting down sink service");
         this.service.shutDown();
         if (server != null) {
             server.shutdown().awaitTermination(30, TimeUnit.SECONDS);
             // force shutdown if not terminated
             if (!server.isTerminated()) {
+                log.info("server did not terminate in {} seconds. Shutting down forcefully", 30);
                 server.shutdownNow();
             }
         }
+        log.info("server.stop successfully completed");
     }
 
     /**

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

@@ -23,9 +23,11 @@
             .newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2);
 
     private final Sinker sinker;
+    private final CompletableFuture<Void> shutdownSignal;
 
-    public Service(Sinker sinker) {
+    public Service(Sinker sinker, CompletableFuture<Void> shutdownSignal) {
         this.sinker = sinker;
+        this.shutdownSignal = shutdownSignal;
     }
 
     /**
@@ -96,15 +98,17 @@
                         datumStream.writeMessage(constructHandlerDatum(request));
                     }
                 } catch (Exception e) {
-                    log.error("Encountered error in sinkFn - {}", e.getMessage());
-                    responseObserver.onError(e);
+                    log.error("Encountered error in sinkFn onNext - {}", e.getMessage());
+                    shutdownSignal.completeExceptionally(e);
+                    responseObserver.onError(Status.INTERNAL.withDescription(e.getMessage()).asException());
                 }
             }
 
             @Override
             public void onError(Throwable throwable) {
                 log.error("Encountered error in sinkFn - {}", throwable.getMessage());
-                responseObserver.onError(throwable);
+                shutdownSignal.completeExceptionally(throwable);
+                responseObserver.onError(Status.INTERNAL.withDescription(throwable.getMessage()).asException());
             }
 
             @Override