feat: use gRPC Bidirectional Streaming for Mapper and Transformer (#143)

Signed-off-by: Yashash H L <[email protected]>

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

@@ -48,7 +48,7 @@ public StreamObserver<Batchmap.BatchMapRequest> batchMapFn(StreamObserver<Batchm
         Future<BatchResponses> result = batchMapTaskExecutor.submit(() -> this.batchMapper.processMessage(
                 datumStream));
 
-        return new StreamObserver<Batchmap.BatchMapRequest>() {
+        return new StreamObserver<>() {
             @Override
             public void onNext(Batchmap.BatchMapRequest mapRequest) {
                 try {

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

@@ -14,4 +14,6 @@ class Constants {
     public static final String MAP_MODE_KEY = "MAP_MODE";
 
     public static final String MAP_MODE = "unary-map";
+
+    public static final String EOF = "EOF";
 }

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

@@ -0,0 +1,142 @@
+package io.numaproj.numaflow.mapper;
+
+import akka.actor.AbstractActor;
+import akka.actor.ActorRef;
+import akka.actor.AllDeadLetters;
+import akka.actor.AllForOneStrategy;
+import akka.actor.Props;
+import akka.actor.SupervisorStrategy;
+import akka.japi.pf.DeciderBuilder;
+import akka.japi.pf.ReceiveBuilder;
+import io.grpc.Status;
+import io.grpc.stub.StreamObserver;
+import io.numaproj.numaflow.map.v1.MapOuterClass;
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * MapSupervisorActor actor is responsible for distributing the messages to actors and handling failure.
+ * It creates a MapperActor for each incoming request and listens to the responses from the MapperActor.
+ * <p>
+ * MapSupervisorActor
+ * │
+ * ├── Creates MapperActor instances for each incoming MapRequest
+ * │   │
+ * │   ├── MapperActor 1
+ * │   │   ├── Processes MapRequest
+ * │   │   ├── Sends MapResponse to MapSupervisorActor
+ * │   │   └── Stops itself after processing
+ * │   │
+ * │   ├── MapperActor 2
+ * │   │   ├── Processes MapRequest
+ * │   │   ├── Sends MapResponse to MapSupervisorActor
+ * │   │   └── Stops itself after processing
+ * │   │
+ * ├── Listens to the responses from the MapperActor instances
+ * │   ├── On receiving a MapResponse, writes the response back to the client
+ * │
+ * ├── If any MapperActor fails (throws an exception):
+ * │   ├── Sends the exception back to the client
+ * │   ├── Initiates a shutdown by completing the CompletableFuture exceptionally
+ * │   └── Stops all child actors (AllForOneStrategy)
+ */
+@Slf4j
+class MapSupervisorActor extends AbstractActor {
+    private final Mapper mapper;
+    private final StreamObserver<MapOuterClass.MapResponse> responseObserver;
+    private final CompletableFuture<Void> failureFuture;
+
+    public MapSupervisorActor(
+            Mapper mapper,
+            StreamObserver<MapOuterClass.MapResponse> responseObserver,
+            CompletableFuture<Void> failureFuture) {
+        this.mapper = mapper;
+        this.responseObserver = responseObserver;
+        this.failureFuture = failureFuture;
+    }
+
+    public static Props props(
+            Mapper mapper,
+            StreamObserver<MapOuterClass.MapResponse> responseObserver,
+            CompletableFuture<Void> failureFuture) {
+        return Props.create(MapSupervisorActor.class, mapper, responseObserver, failureFuture);
+    }
+
+    @Override
+    public void preRestart(Throwable reason, Optional<Object> message) {
+        log.debug("supervisor pre restart was executed");
+        failureFuture.completeExceptionally(reason);
+        responseObserver.onError(Status.UNKNOWN
+                .withDescription(reason.getMessage())
+                .withCause(reason)
+                .asException());
+        Service.mapperActorSystem.stop(getSelf());
+    }
+
+    @Override
+    public void postStop() {
+        log.debug("post stop of supervisor executed - {}", getSelf().toString());
+    }
+
+    @Override
+    public Receive createReceive() {
+        return ReceiveBuilder
+                .create()
+                .match(MapOuterClass.MapRequest.class, this::processRequest)
+                .match(MapOuterClass.MapResponse.class, this::sendResponse)
+                .match(Exception.class, this::handleFailure)
+                .match(AllDeadLetters.class, this::handleDeadLetters)
+                .match(String.class, eof -> responseObserver.onCompleted())
+                .build();
+    }
+
+    private void handleFailure(Exception e) {
+        responseObserver.onError(Status.UNKNOWN
+                .withDescription(e.getMessage())
+                .withCause(e)
+                .asException());
+        failureFuture.completeExceptionally(e);
+    }
+
+    private void sendResponse(MapOuterClass.MapResponse mapResponse) {
+        responseObserver.onNext(mapResponse);
+    }
+
+    private void processRequest(MapOuterClass.MapRequest mapRequest) {
+        // Create a MapperActor for each incoming request.
+        ActorRef mapperActor = getContext()
+                .actorOf(MapperActor.props(
+                        mapper));
+
+        // Send the message to the MapperActor.
+        mapperActor.tell(mapRequest, getSelf());
+    }
+
+    // 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");
+        responseObserver.onError(Status.UNKNOWN.withDescription("dead letters").asException());
+        failureFuture.completeExceptionally(new Throwable("dead letters"));
+        getContext().getSystem().stop(getSelf());
+    }
+
+    @Override
+    public SupervisorStrategy supervisorStrategy() {
+        // we want to stop all child actors in case of any exception
+        return new AllForOneStrategy(
+                DeciderBuilder
+                        .match(Exception.class, e -> {
+                            failureFuture.completeExceptionally(e);
+                            responseObserver.onError(Status.UNKNOWN
+                                    .withDescription(e.getMessage())
+                                    .withCause(e)
+                                    .asException());
+                            return SupervisorStrategy.stop();
+                        })
+                        .build()
+        );
+    }
+}

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

@@ -0,0 +1,87 @@
+package io.numaproj.numaflow.mapper;
+
+import akka.actor.AbstractActor;
+import akka.actor.Props;
+import akka.japi.pf.ReceiveBuilder;
+import com.google.protobuf.ByteString;
+import io.numaproj.numaflow.map.v1.MapOuterClass;
+
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Mapper actor that processes the map request. It invokes the mapper to process the request and
+ * sends the response back to the sender actor(MapSupervisorActor). In case of any exception, it
+ * sends the exception back to the sender actor. It stops itself after processing the request.
+ */
+class MapperActor extends AbstractActor {
+    private final Mapper mapper;
+
+    public MapperActor(Mapper mapper) {
+        this.mapper = mapper;
+    }
+
+    public static Props props(Mapper mapper) {
+        return Props.create(MapperActor.class, mapper);
+    }
+
+    @Override
+    public Receive createReceive() {
+        return ReceiveBuilder.create()
+                .match(MapOuterClass.MapRequest.class, this::processRequest)
+                .build();
+    }
+
+    /**
+     * Process the map request and send the response back to the sender actor.
+     *
+     * @param mapRequest map request
+     */
+    private void processRequest(MapOuterClass.MapRequest mapRequest) {
+        Datum handlerDatum = new HandlerDatum(
+                mapRequest.getRequest().getValue().toByteArray(),
+                Instant.ofEpochSecond(
+                        mapRequest.getRequest().getWatermark().getSeconds(),
+                        mapRequest.getRequest().getWatermark().getNanos()),
+                Instant.ofEpochSecond(
+                        mapRequest.getRequest().getEventTime().getSeconds(),
+                        mapRequest.getRequest().getEventTime().getNanos()),
+                mapRequest.getRequest().getHeadersMap()
+        );
+        String[] keys = mapRequest.getRequest().getKeysList().toArray(new String[0]);
+        try {
+            MessageList resultMessages = this.mapper.processMessage(keys, handlerDatum);
+            MapOuterClass.MapResponse response = buildResponse(resultMessages, mapRequest.getId());
+            getSender().tell(response, getSelf());
+        } catch (Exception e) {
+            getSender().tell(e, getSelf());
+        }
+        context().stop(getSelf());
+    }
+
+    /**
+     * Build the response from the message list.
+     *
+     * @param messageList message list
+     *
+     * @return map response
+     */
+    private MapOuterClass.MapResponse buildResponse(MessageList messageList, String ID) {
+        MapOuterClass.MapResponse.Builder responseBuilder = MapOuterClass
+                .MapResponse
+                .newBuilder();
+
+        messageList.getMessages().forEach(message -> {
+            responseBuilder.addResults(MapOuterClass.MapResponse.Result.newBuilder()
+                    .setValue(message.getValue() == null ? ByteString.EMPTY : ByteString.copyFrom(
+                            message.getValue()))
+                    .addAllKeys(message.getKeys()
+                            == null ? new ArrayList<>() : List.of(message.getKeys()))
+                    .addAllTags(message.getTags()
+                            == null ? new ArrayList<>() : List.of(message.getTags()))
+                    .build());
+        });
+        return responseBuilder.setId(ID).build();
+    }
+}