reducers

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

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

@@ -6,6 +6,7 @@
 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;
 
@@ -21,6 +22,7 @@ public class Server {
     private final Service service;
     private final ServerInfoAccessor serverInfoAccessor = new ServerInfoAccessorImpl(new ObjectMapper());
     private io.grpc.Server server;
+    private final GrpcServerHelper grpcServerHelper;
 
     /**
      * constructor to create gRPC server.
@@ -40,6 +42,7 @@ public Server(ReducerFactory<? extends Reducer> reducerFactory) {
     public Server(ReducerFactory<? extends Reducer> reducerFactory, GRPCConfig grpcConfig) {
         this.service = new Service(reducerFactory);
         this.grpcConfig = grpcConfig;
+        this.grpcServerHelper = new GrpcServerHelper();
     }
 
     /**
@@ -57,33 +60,32 @@ public void start() throws Exception {
         }
 
         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");
+            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);
@@ -99,7 +101,9 @@ public void start() throws Exception {
      * @throws InterruptedException if the current thread is interrupted while waiting
      */
     public void awaitTermination() throws InterruptedException {
+        log.info("reducer server is waiting for termination");
         server.awaitTermination();
+        log.info("reducer server has terminated");
     }
 
     /**

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

@@ -8,6 +8,7 @@
 import io.numaproj.numaflow.info.ServerInfoAccessorImpl;
 import io.numaproj.numaflow.reducestreamer.model.ReduceStreamer;
 import io.numaproj.numaflow.reducestreamer.model.ReduceStreamerFactory;
+import io.numaproj.numaflow.shared.GrpcServerHelper;
 import io.numaproj.numaflow.shared.GrpcServerUtils;
 import lombok.extern.slf4j.Slf4j;
 
@@ -22,6 +23,7 @@ public class Server {
     private final Service service;
     private final ServerInfoAccessor serverInfoAccessor = new ServerInfoAccessorImpl(new ObjectMapper());
     private io.grpc.Server server;
+    private final GrpcServerHelper grpcServerHelper;
 
     /**
      * constructor to create gRPC server.
@@ -43,6 +45,7 @@ public Server(
             GRPCConfig grpcConfig) {
         this.service = new Service(reduceStreamerFactory);
         this.grpcConfig = grpcConfig;
+        this.grpcServerHelper = new GrpcServerHelper();
     }
 
     /**
@@ -60,33 +63,32 @@ public void start() throws Exception {
         }
 
         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");
+            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);
@@ -102,7 +104,9 @@ public void start() throws Exception {
      * @throws InterruptedException if the current thread is interrupted while waiting
      */
     public void awaitTermination() throws InterruptedException {
+        log.info("reduce stream server is waiting for termination");
         server.awaitTermination();
+        log.info("reduce stream server terminated");
     }
 
     /**

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

@@ -8,6 +8,7 @@
 import io.numaproj.numaflow.info.ServerInfoAccessorImpl;
 import io.numaproj.numaflow.sessionreducer.model.SessionReducer;
 import io.numaproj.numaflow.sessionreducer.model.SessionReducerFactory;
+import io.numaproj.numaflow.shared.GrpcServerHelper;
 import io.numaproj.numaflow.shared.GrpcServerUtils;
 import lombok.extern.slf4j.Slf4j;
 
@@ -22,6 +23,7 @@ public class Server {
     private final Service service;
     private final ServerInfoAccessor serverInfoAccessor = new ServerInfoAccessorImpl(new ObjectMapper());
     private io.grpc.Server server;
+    private final GrpcServerHelper grpcServerHelper;
 
     /**
      * constructor to create gRPC server.
@@ -43,6 +45,7 @@ public Server(
             GRPCConfig grpcConfig) {
         this.service = new Service(sessionReducerFactory);
         this.grpcConfig = grpcConfig;
+        this.grpcServerHelper = new GrpcServerHelper();
     }
 
     /**
@@ -60,33 +63,32 @@ public void start() throws Exception {
         }
 
         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");
+            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);
@@ -102,7 +104,9 @@ public void start() throws Exception {
      * @throws InterruptedException if the current thread is interrupted while waiting
      */
     public void awaitTermination() throws InterruptedException {
+        log.info("session reduce server is waiting for termination");
         server.awaitTermination();
+        log.info("session reduce server terminated");
     }
 
     /**

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

@@ -127,73 +127,4 @@ public static void writeServerInfo(
         log.info("Writing server info {} to {}", serverInfo, infoFilePath);
         serverInfoAccessor.write(serverInfo, infoFilePath);
     }
-
-    // TODO - remove this one
-    public static ServerBuilder<?> createServerBuilder(
-            String socketPath,
-            int maxMessageSize,
-            boolean isLocal,
-            int port) {
-        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);
-        }
-
-        return NettyServerBuilder
-                .forAddress(new DomainSocketAddress(socketPath))
-                .channelType(GrpcServerUtils.getChannelTypeClass())
-                .maxInboundMessageSize(maxMessageSize)
-                .bossEventLoopGroup(GrpcServerUtils.createEventLoopGroup(1, "netty-boss"))
-                .workerEventLoopGroup(GrpcServerUtils.createEventLoopGroup(
-                        ThreadUtils.INSTANCE.availableProcessors(),
-                        "netty-worker"))
-                .intercept(interceptor);
-    }
 }

src/test/java/io/numaproj/numaflow/shared/GrpcServerUtilsTest.java

@@ -36,16 +36,6 @@ public void testWriteServerInfo() throws Exception {
                 .write(Mockito.any(), Mockito.eq("infoFilePath"));
     }
 
-    @Test
-    public void testCreateServerBuilder() {
-        ServerBuilder<?> serverBuilder = GrpcServerUtils.createServerBuilder(
-                "socketPath",
-                1000,
-                false,
-                50051);
-        Assert.assertNotNull(serverBuilder);
-    }
-
     @Test
     public void testWindowStartTime() {
         Context.Key<String> windowStartTime = GrpcServerUtils.WINDOW_START_TIME;