fix: remove System.exit from JVM shutdown hooks

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

@@ -12,13 +12,20 @@
 
 import java.util.Collections;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 /**
  * Server is the gRPC server for executing map operation.
  */
 @Slf4j
 public class Server {
 
+    // Upper bound on how long a graceful shutdown waits for the actor system to
+    // terminate before giving up, so a stuck actor cannot hang shutdown forever.
+    private static final long ACTOR_SYSTEM_TERMINATION_TIMEOUT_SECONDS = 30;
+
     private final GRPCConfig grpcConfig;
     private final CompletableFuture<Void> shutdownSignal;
     private final ServerInfoAccessor serverInfoAccessor = new ServerInfoAccessorImpl(new ObjectMapper());
@@ -76,10 +83,13 @@ public void start() throws Exception {
                 // 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");
                 try {
+                    // NOTE: never call System.exit from a shutdown hook - the JVM is already
+                    // shutting down and Shutdown.exit blocks forever on the lock held by the
+                    // thread running the hooks, deadlocking the process.
                     this.stop();
-                    // FIXME - this is a workaround to immediately terminate the JVM process
-                    // The correct way to do this is to stop all the actors and wait for them to terminate
-                    System.exit(0);
+                    // Stop all actors and wait for them to terminate so shutdown is graceful.
+                    // The JVM exits on its own once the hook returns.
+                    shutdownActorSystem();
                 } catch (InterruptedException e) {
                     Thread.interrupted();
                     e.printStackTrace(System.err);
@@ -100,8 +110,10 @@ public void start() throws Exception {
                 System.err.println("*** shutting down mapper gRPC server because of an exception - " + e.getMessage());
                 try {
                     this.stop();
-                    // FIXME - this is a workaround to immediately terminate the JVM process
-                    // The correct way to do this is to stop all the actors and wait for them to terminate
+                    // Force the process to exit so the platform restarts the container - one
+                    // unrecoverable user error should trigger a restart. System.exit is safe here
+                    // (this is not a shutdown hook) and it runs the shutdown hook, which
+                    // gracefully terminates the actor system before the JVM dies.
                     System.exit(0);
                 } catch (InterruptedException ex) {
                     Thread.interrupted();
@@ -133,4 +145,26 @@ public void awaitTermination() throws InterruptedException {
     public void stop() throws InterruptedException {
         this.server.gracefullyShutdown();
     }
+
+    /**
+     * Terminates the mapper actor system and waits, bounded by
+     * {@link #ACTOR_SYSTEM_TERMINATION_TIMEOUT_SECONDS}, for all actors to finish. This makes
+     * shutdown graceful instead of relying on an abrupt process kill. Failures are only logged
+     * (to stderr, since the logger may already be torn down during JVM shutdown) so that shutdown
+     * always proceeds.
+     */
+    private void shutdownActorSystem() {
+        try {
+            Service.mapperActorSystem.terminate();
+            Service.mapperActorSystem
+                    .getWhenTerminated()
+                    .toCompletableFuture()
+                    .get(ACTOR_SYSTEM_TERMINATION_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            System.err.println("*** interrupted while waiting for mapper actor system to terminate");
+        } catch (ExecutionException | TimeoutException e) {
+            System.err.println("*** mapper actor system did not terminate cleanly - " + e.getMessage());
+        }
+    }
 }

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

@@ -11,13 +11,20 @@
 import lombok.extern.slf4j.Slf4j;
 
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 /**
  * Server is the gRPC server for executing source transformer operation.
  */
 @Slf4j
 public class Server {
 
+    // Upper bound on how long a graceful shutdown waits for the actor system to
+    // terminate before giving up, so a stuck actor cannot hang shutdown forever.
+    private static final long ACTOR_SYSTEM_TERMINATION_TIMEOUT_SECONDS = 30;
+
     private final GRPCConfig grpcConfig;
     private final CompletableFuture<Void> shutdownSignal;
     private final ServerInfoAccessor serverInfoAccessor = new ServerInfoAccessorImpl(new ObjectMapper());
@@ -72,10 +79,13 @@ public void start() throws Exception {
                 // 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");
                 try {
+                    // NOTE: never call System.exit from a shutdown hook - the JVM is already
+                    // shutting down and Shutdown.exit blocks forever on the lock held by the
+                    // thread running the hooks, deadlocking the process.
                     this.stop();
-                    // FIXME - this is a workaround to immediately terminate the JVM process
-                    // The correct way to do this is to stop all the actors and wait for them to terminate
-                    System.exit(0);
+                    // Stop all actors and wait for them to terminate so shutdown is graceful.
+                    // The JVM exits on its own once the hook returns.
+                    shutdownActorSystem();
                 } catch (InterruptedException e) {
                     Thread.interrupted();
                     e.printStackTrace(System.err);
@@ -96,8 +106,10 @@ public void start() throws Exception {
                 System.err.println("*** shutting down transformer gRPC server because of an exception - " + e.getMessage());
                 try {
                     this.stop();
-                    // FIXME - this is a workaround to immediately terminate the JVM process
-                    // The correct way to do this is to stop all the actors and wait for them to terminate
+                    // Force the process to exit so the platform restarts the container - one
+                    // unrecoverable user error should trigger a restart. System.exit is safe here
+                    // (this is not a shutdown hook) and it runs the shutdown hook, which
+                    // gracefully terminates the actor system before the JVM dies.
                     System.exit(0);
                 } catch (InterruptedException ex) {
                     Thread.interrupted();
@@ -129,4 +141,26 @@ public void awaitTermination() throws InterruptedException {
     public void stop() throws InterruptedException {
         this.server.gracefullyShutdown();
     }
+
+    /**
+     * Terminates the transformer actor system and waits, bounded by
+     * {@link #ACTOR_SYSTEM_TERMINATION_TIMEOUT_SECONDS}, for all actors to finish. This makes
+     * shutdown graceful instead of relying on an abrupt process kill. Failures are only logged
+     * (to stderr, since the logger may already be torn down during JVM shutdown) so that shutdown
+     * always proceeds.
+     */
+    private void shutdownActorSystem() {
+        try {
+            Service.transformerActorSystem.terminate();
+            Service.transformerActorSystem
+                    .getWhenTerminated()
+                    .toCompletableFuture()
+                    .get(ACTOR_SYSTEM_TERMINATION_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            System.err.println("*** interrupted while waiting for transformer actor system to terminate");
+        } catch (ExecutionException | TimeoutException e) {
+            System.err.println("*** transformer actor system did not terminate cleanly - " + e.getMessage());
+        }
+    }
 }

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

@@ -0,0 +1,98 @@
+package io.numaproj.numaflow.mapper;
+
+import org.junit.Test;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.InputStreamReader;
+import java.nio.charset.StandardCharsets;
+import java.util.concurrent.TimeUnit;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+/**
+ * Reproduces the shutdown-hook self-deadlock fixed by removing {@code System.exit(0)}
+ * from the JVM shutdown hook in {@link Server}.
+ * <p>
+ * The deadlock only manifests at JVM exit, so it cannot be observed from a normal
+ * in-process test method. Instead we fork a child JVM ({@link ShutdownHookReproMain})
+ * that starts a real {@link Server} (which registers the hook) and then initiates
+ * shutdown, and we assert the process actually terminates within a timeout.
+ * <p>
+ * Expected results:
+ * <ul>
+ *   <li>buggy code (hook calls {@code System.exit(0)}): child JVM hangs -> this test FAILS.</li>
+ *   <li>fixed code (hook only calls {@code stop()}): child JVM exits 0 -> this test PASSES.</li>
+ * </ul>
+ */
+public class ServerShutdownHookTest {
+
+    private static final long EXIT_TIMEOUT_SECONDS = 20;
+
+    @Test
+    public void shutdownHookMustNotDeadlockJvmExit() throws Exception {
+        String javaBin = System.getProperty("java.home")
+                + File.separator + "bin" + File.separator + "java";
+        String classpath = System.getProperty("java.class.path");
+
+        ProcessBuilder pb = new ProcessBuilder(
+                javaBin,
+                "-cp", classpath,
+                ShutdownHookReproMain.class.getName());
+        pb.redirectErrorStream(true);
+
+        Process process = pb.start();
+
+        // Drain the child's output on a background thread so its pipe buffer can never
+        // fill up and block the child (which would be a false-positive "hang").
+        StringBuilder output = new StringBuilder();
+        Thread drainer = new Thread(() -> {
+            try (BufferedReader reader = new BufferedReader(new InputStreamReader(
+                    process.getInputStream(), StandardCharsets.UTF_8))) {
+                String line;
+                while ((line = reader.readLine()) != null) {
+                    synchronized (output) {
+                        output.append(line).append('\n');
+                    }
+                }
+            } catch (Exception ignored) {
+                // stream closed on process exit
+            }
+        });
+        drainer.setDaemon(true);
+        drainer.start();
+
+        boolean exited = process.waitFor(EXIT_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+
+        if (!exited) {
+            process.destroyForcibly();
+            process.waitFor(10, TimeUnit.SECONDS);
+            drainer.join(TimeUnit.SECONDS.toMillis(2));
+            synchronized (output) {
+                fail("Forked JVM failed to exit within " + EXIT_TIMEOUT_SECONDS
+                        + "s: the shutdown hook self-deadlocked (System.exit() called "
+                        + "from within the hook). Child output:\n" + output);
+            }
+        }
+
+        drainer.join(TimeUnit.SECONDS.toMillis(2));
+        String childOutput;
+        synchronized (output) {
+            childOutput = output.toString();
+        }
+
+        // Make sure the child actually got far enough to register the hook; otherwise a
+        // clean exit would be meaningless (e.g. an unrelated startup failure).
+        assertTrue(
+                "Child JVM did not report a successful server start; cannot trust the exit "
+                        + "result. Child output:\n" + childOutput,
+                childOutput.contains("SERVER_STARTED"));
+
+        assertEquals(
+                "Forked JVM should exit cleanly (exit code 0). Child output:\n" + childOutput,
+                0,
+                process.exitValue());
+    }
+}

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

@@ -0,0 +1,63 @@
+package io.numaproj.numaflow.mapper;
+
+import io.grpc.inprocess.InProcessServerBuilder;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+/**
+ * Helper entry point launched in a *forked* JVM by {@link ServerShutdownHookTest}.
+ * <p>
+ * It starts a real {@link Server} with {@code isLocal=false} so that
+ * {@link Server#start()} registers the JVM shutdown hook, then triggers a normal
+ * JVM shutdown from the main thread (exactly like the surefire booter does at the
+ * end of a test run).
+ * <p>
+ * With the buggy hook - which calls {@code System.exit(0)} from *inside* the hook -
+ * the JVM self-deadlocks: the thread running the hooks holds the {@code Shutdown}
+ * lock and waits for the hook to finish, while the hook blocks in {@code Shutdown.exit}
+ * waiting for that same lock. The process then hangs forever. With the fix (hook only
+ * calls {@code stop()} and terminates the actor system) the JVM exits cleanly with code 0.
+ */
+public class ShutdownHookReproMain {
+
+    public static void main(String[] args) throws Exception {
+        Path tmpDir = Files.createTempDirectory("map-shutdown-repro");
+        String socketPath = tmpDir.resolve("map.sock").toString();
+        String infoFilePath = tmpDir.resolve("server-info").toString();
+
+        GRPCConfig config = GRPCConfig.newBuilder()
+                .maxMessageSize(Constants.DEFAULT_MESSAGE_SIZE)
+                .socketPath(socketPath)
+                .infoFilePath(infoFilePath)
+                // isLocal=false is what makes Server.start() register the shutdown hook.
+                .isLocal(false)
+                .build();
+
+        // Use the in-process (VisibleForTesting) server so we don't need to bind a
+        // real unix-domain socket; the shutdown hook is registered regardless.
+        Server server = new Server(
+                config,
+                new NoopMapper(),
+                null,
+                InProcessServerBuilder.generateName());
+
+        server.start();
+
+        // Signal to the parent that startup succeeded, so a failure to exit can be
+        // attributed to the shutdown hook rather than a startup problem.
+        System.out.println("SERVER_STARTED");
+        System.out.flush();
+
+        // Initiate a normal JVM shutdown from the main thread. This fires the hook
+        // registered by Server.start(). If the hook calls System.exit(), we deadlock.
+        System.exit(0);
+    }
+
+    private static final class NoopMapper extends Mapper {
+        @Override
+        public MessageList processMessage(String[] keys, Datum datum) {
+            return MessageList.newBuilder().build();
+        }
+    }
+}