akka change for mapper

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

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

@@ -47,6 +47,8 @@ class MapSupervisorActor extends AbstractActor {
     private final Mapper mapper;
     private final StreamObserver<MapOuterClass.MapResponse> responseObserver;
     private final CompletableFuture<Void> shutdownSignal;
+    private int activeMapperCount;
+    private Exception userException;
 
     public MapSupervisorActor(
             Mapper mapper,
@@ -55,6 +57,8 @@ public MapSupervisorActor(
         this.mapper = mapper;
         this.responseObserver = responseObserver;
         this.shutdownSignal = failureFuture;
+        this.userException = null;
+        this.activeMapperCount = 0;
     }
 
     public static Props props(
@@ -66,13 +70,14 @@ public static Props props(
 
     @Override
     public void preRestart(Throwable reason, Optional<Object> message) {
-        log.debug("supervisor pre restart was executed");
+        getContext().getSystem().log().warning("supervisor pre restart was executed due to: {}", reason.getMessage());
         shutdownSignal.completeExceptionally(reason);
         responseObserver.onError(Status.INTERNAL
                 .withDescription(reason.getMessage())
                 .withCause(reason)
                 .asException());
         Service.mapperActorSystem.stop(getSelf());
+        shutdownSignal.completeExceptionally(reason);
     }
 
     @Override
@@ -94,34 +99,52 @@ public Receive createReceive() {
 
     private void handleFailure(Exception e) {
         log.error("Encountered error in mapFn - {}", e.getMessage());
-        shutdownSignal.completeExceptionally(e);
-        responseObserver.onError(Status.INTERNAL
-                .withDescription(e.getMessage())
-                .withCause(e)
-                .asException());
+        if (userException == null) {
+            userException = e;
+            // only send the very first exception to the client
+            // one exception should trigger a container restart
+            responseObserver.onError(Status.INTERNAL
+                    .withDescription(e.getMessage())
+                    .withCause(e)
+                    .asException());
+        }
+        activeMapperCount--;
     }
 
     private void sendResponse(MapOuterClass.MapResponse mapResponse) {
         responseObserver.onNext(mapResponse);
+        activeMapperCount--;
     }
 
     private void processRequest(MapOuterClass.MapRequest mapRequest) {
+        if (userException != null) {
+            log.info("a previous mapper actor failed, not processing any more requests");
+            if (activeMapperCount == 0) {
+                log.info("there is no more active mapper AKKA actors - stopping the system");
+                getContext().getSystem().stop(getSelf());
+                log.info("AKKA system stopped");
+                shutdownSignal.completeExceptionally(userException);
+            }
+            return;
+        }
+
         // Create a MapperActor for each incoming request.
         ActorRef mapperActor = getContext()
                 .actorOf(MapperActor.props(
                         mapper));
 
         // Send the message to the MapperActor.
         mapperActor.tell(mapRequest, getSelf());
+        activeMapperCount++;
     }
 
     // if we see dead letters, we need to stop the execution and exit
     // to make sure no messages are lost
     private void handleDeadLetters(AllDeadLetters deadLetter) {
         log.debug("got a dead letter, stopping the execution");
-        shutdownSignal.completeExceptionally(new Throwable("dead letters"));
         responseObserver.onError(Status.INTERNAL.withDescription("dead letters").asException());
         getContext().getSystem().stop(getSelf());
+        shutdownSignal.completeExceptionally(new Throwable("dead letters"));
     }
 
     @Override

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

@@ -93,6 +93,9 @@ public void start() throws Exception {
                 Server.this.stop();
                 log.info("gracefully shutting down event loop groups");
                 this.grpcServerHelper.gracefullyShutdownEventLoopGroups();
+                // 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);
             } catch (InterruptedException e) {
                 Thread.interrupted();
                 e.printStackTrace(System.err);
@@ -112,6 +115,9 @@ public void start() throws Exception {
                     Server.this.stop();
                     log.info("gracefully shutting down event loop groups");
                     this.grpcServerHelper.gracefullyShutdownEventLoopGroups();
+                    // 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);
                 } catch (InterruptedException ex) {
                     Thread.interrupted();
                     ex.printStackTrace(System.err);

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

@@ -191,7 +191,7 @@ private void processRequest(Sourcetransformer.SourceTransformRequest transformRe
      * @param deadLetter The dead letter to be handled.
      */
     private void handleDeadLetters(AllDeadLetters deadLetter) {
-        log.info("got a dead letter, stopping the execution");
+        log.debug("got a dead letter, stopping the execution");
         responseObserver.onError(Status.INTERNAL.withDescription("dead letters").asException());
         getContext().getSystem().stop(getSelf());
         shutdownSignal.completeExceptionally(new Throwable("dead letters"));