feat: bidirectional streaming (#23)

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

Author: yhl25

src/main/java/io/numaproj/numaflow/function/FunctionService.java

@@ -164,16 +164,16 @@ public StreamObserver<Udfunction.Datum> reduceFn(final StreamObserver<Udfunction
             create a supervisor actor which assign the tasks to child actors.
             we create a child actor for every key in a window.
         */
-        ActorRef parentActorRef = actorSystem
-                .actorOf(ReduceSupervisorActor.props(reducerFactory, md, shutdownActorRef));
+        ActorRef supervisorActor = actorSystem
+                .actorOf(ReduceSupervisorActor.props(reducerFactory, md, shutdownActorRef, responseObserver));
 
 
         return new StreamObserver<Udfunction.Datum>() {
             @Override
             public void onNext(Udfunction.Datum datum) {
                 // send the message to parent actor, which takes care of distribution.
-                if (!parentActorRef.isTerminated()) {
-                    parentActorRef.tell(datum, parentActorRef);
+                if (!supervisorActor.isTerminated()) {
+                    supervisorActor.tell(datum, supervisorActor);
                 } else {
                     responseObserver.onError(new Throwable("Supervisor actor was terminated"));
                 }
@@ -187,21 +187,8 @@ public void onError(Throwable throwable) {
 
             @Override
             public void onCompleted() {
-
-                // Ask the parent to return the list of futures returned by child actors.
-                Future<Object> resultFuture = Patterns.ask(parentActorRef, EOF, Integer.MAX_VALUE);
-
-                List<Future<Object>> udfResultFutures;
-                try {
-                    udfResultFutures = (List<Future<Object>>) Await.result(
-                            resultFuture,
-                            Duration.Inf());
-                } catch (TimeoutException | InterruptedException e) {
-                    responseObserver.onError(e);
-                    return;
-                }
-
-                extractResult(udfResultFutures, responseObserver, parentActorRef);
+                // indicate the end of input to the supervisor
+                supervisorActor.tell(EOF, ActorRef.noSender());
 
             }
         };
@@ -243,52 +230,6 @@ private Udfunction.DatumList buildDatumListResponse(MessageT[] messageTs) {
         return datumListBuilder.build();
     }
 
-    /*
-        extracts and returns the result to the observer.
-     */
-    private void extractResult(
-            List<Future<Object>> futureList,
-            StreamObserver<Udfunction.DatumList> responseObserver,
-            ActorRef supervisorActorRef) {
-        Udfunction.DatumList.Builder responseBuilder = Udfunction.DatumList.newBuilder();
-
-        // build the response when its completed.
-        Futures
-                .sequence(futureList, actorSystem.dispatcher())
-                .onComplete(new OnComplete<>() {
-                    @Override
-                    public void onComplete(
-                            Throwable failure,
-                            Iterable<Object> success) {
-
-                        // if there are any failures indicate it to the observer.
-                        if (failure != null) {
-                            log.error("Error while getting output from actors - {}"
-                                    , failure.getMessage());
-
-                            responseObserver.onError(failure);
-                            return;
-                        }
-
-                        success.forEach(ele -> {
-                            Udfunction.DatumList list = buildDatumListResponse((Message[]) ele);
-                            responseBuilder.addAllElements(list.getElementsList());
-                        });
-                        Udfunction.DatumList response = responseBuilder.build();
-
-                        responseObserver.onNext(response);
-                        responseObserver.onCompleted();
-
-                        /*
-                            once the result is returned we can stop the supervisor actor.
-                            stopping the supervisor will stop all its child actors.
-                            we should  explicitly stop the actors for it to be garbage collected.
-                        */
-                        actorSystem.stop(supervisorActorRef);
-                    }
-                }, actorSystem.dispatcher());
-    }
-
     // log the exception and exit if there are any uncaught exceptions.
     private void handleFailure(CompletableFuture<Void> failureFuture) {
         new Thread(() -> {

src/main/java/io/numaproj/numaflow/function/reduce/ActorResponse.java

@@ -0,0 +1,14 @@
+package io.numaproj.numaflow.function.reduce;
+
+import io.numaproj.numaflow.function.v1.Udfunction;
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+/*
+    used to store the reduced result from the handler
+*/
+@Getter
+@AllArgsConstructor
+public class ActorResponse {
+    String key;
+    Udfunction.DatumList datumList;
+}

src/main/java/io/numaproj/numaflow/function/reduce/ReduceActor.java

@@ -3,19 +3,22 @@
 import akka.actor.AbstractActor;
 import akka.actor.Props;
 import akka.japi.pf.ReceiveBuilder;
+import com.google.protobuf.ByteString;
 import io.numaproj.numaflow.function.HandlerDatum;
+import io.numaproj.numaflow.function.Message;
 import io.numaproj.numaflow.function.metadata.Metadata;
+import io.numaproj.numaflow.function.v1.Udfunction;
 import lombok.AllArgsConstructor;
-import lombok.NoArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
+import java.util.Arrays;
+
 /**
  * Reduce actor invokes the user defined code and returns the result.
  */
 
 @Slf4j
 @AllArgsConstructor
-@NoArgsConstructor
 public class ReduceActor extends AbstractActor {
 
     private String key;
@@ -40,7 +43,20 @@ private void invokeHandler(HandlerDatum handlerDatum) {
     }
 
     private void getResult(String eof) {
-        getSender().tell(this.groupBy.getOutput(key, md), getSelf());
+        Message[] resultMessages = this.groupBy.getOutput(key, md);
+        // send the result back to sender(parent actor)
+        getSender().tell(buildDatumListResponse(resultMessages), getSelf());
+    }
+
+    private ActorResponse buildDatumListResponse(Message[] messages) {
+        Udfunction.DatumList.Builder datumListBuilder = Udfunction.DatumList.newBuilder();
+        Arrays.stream(messages).forEach(message -> {
+            datumListBuilder.addElements(Udfunction.Datum.newBuilder()
+                    .setKey(message.getKey())
+                    .setValue(ByteString.copyFrom(message.getValue()))
+                    .build());
+        });
+        return new ActorResponse(this.key, datumListBuilder.build());
     }
 
 }

src/main/java/io/numaproj/numaflow/function/reduce/ReduceSupervisorActor.java

@@ -7,8 +7,8 @@
 import akka.actor.SupervisorStrategy;
 import akka.japi.pf.DeciderBuilder;
 import akka.japi.pf.ReceiveBuilder;
-import akka.pattern.Patterns;
 import com.google.common.base.Preconditions;
+import io.grpc.stub.StreamObserver;
 import io.numaproj.numaflow.function.Function;
 import io.numaproj.numaflow.function.FunctionService;
 import io.numaproj.numaflow.function.HandlerDatum;
@@ -19,7 +19,6 @@
 import scala.collection.Iterable;
 import scala.concurrent.Future;
 
-import java.lang.reflect.InvocationTargetException;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -36,23 +35,27 @@ public class ReduceSupervisorActor extends AbstractActor {
     private final ReducerFactory<? extends Reducer> reducerFactory;
     private final Metadata md;
     private final ActorRef shutdownActor;
+    private final StreamObserver<Udfunction.DatumList> responseObserver;
     private final Map<String, ActorRef> actorsMap = new HashMap<>();
     private final List<Future<Object>> results = new ArrayList<>();
 
     public ReduceSupervisorActor(
             ReducerFactory<? extends Reducer> reducerFactory,
             Metadata md,
-            ActorRef shutdownActor) {
+            ActorRef shutdownActor,
+            StreamObserver<Udfunction.DatumList> responseObserver) {
         this.reducerFactory = reducerFactory;
         this.md = md;
         this.shutdownActor = shutdownActor;
+        this.responseObserver = responseObserver;
     }
 
     public static Props props(
             ReducerFactory<? extends Reducer> reducerFactory,
             Metadata md,
-            ActorRef shutdownActor) {
-        return Props.create(ReduceSupervisorActor.class, reducerFactory, md, shutdownActor);
+            ActorRef shutdownActor,
+            StreamObserver<Udfunction.DatumList> responseObserver) {
+        return Props.create(ReduceSupervisorActor.class, reducerFactory, md, shutdownActor, responseObserver);
     }
 
     // if there is an uncaught exception stop in the supervisor actor, send a signal to shut down
@@ -65,7 +68,7 @@ public void preRestart(Throwable reason, Optional<Object> message) {
 
     @Override
     public SupervisorStrategy supervisorStrategy() {
-        return new ReduceSupervisorStratergy();
+        return new ReduceSupervisorStrategy();
     }
 
 
@@ -81,15 +84,20 @@ public Receive createReceive() {
                 .create()
                 .match(Udfunction.Datum.class, this::invokeActors)
                 .match(String.class, this::sendEOF)
+                .match(ActorResponse.class, this::responseListener)
                 .build();
     }
 
+    /*
+        based on key of the input message invoke the right actor
+        if there is no actor for an incoming key, create a new actor
+        track all the child actors using actors map
+     */
     private void invokeActors(Udfunction.Datum datum) {
         if (!actorsMap.containsKey(datum.getKey())) {
 
             Reducer reducer = reducerFactory.createReducer();
 
-
             ActorRef actorRef = getContext()
                     .actorOf(ReduceActor.props(datum.getKey(), md, reducer));
 
@@ -101,10 +109,25 @@ private void invokeActors(Udfunction.Datum datum) {
 
     private void sendEOF(String EOF) {
         for (Map.Entry<String, ActorRef> entry : actorsMap.entrySet()) {
-            results.add(Patterns.ask(entry.getValue(), EOF, Integer.MAX_VALUE));
+            entry.getValue().tell(EOF, getSelf());
+        }
+    }
+
+    // listen to child actors for the result.
+    private void responseListener(ActorResponse actorResponse) {
+        /*
+            send the result back to the client
+            remove the child entry from the map after getting result.
+            if there are no entries in the map, that means processing is
+            done we can close the stream.
+         */
+
+        responseObserver.onNext(actorResponse.getDatumList());
+        actorsMap.remove(actorResponse.getKey());
+        if (actorsMap.isEmpty()) {
+            responseObserver.onCompleted();
+            getContext().getSystem().stop(getSelf());
         }
-        actorsMap.clear();
-        getSender().tell(results, getSelf());
     }
 
     private HandlerDatum constructHandlerDatum(Udfunction.Datum datum) {
@@ -123,7 +146,7 @@ private HandlerDatum constructHandlerDatum(Udfunction.Datum datum) {
         actors will be restarted, but we want to escalate the exception and terminate
         the system.
     */
-    private final class ReduceSupervisorStratergy extends SupervisorStrategy {
+    private final class ReduceSupervisorStrategy extends SupervisorStrategy {
 
         @Override
         public PartialFunction<Throwable, Directive> decider() {

src/main/proto/function/v1/udfunction.proto

@@ -17,7 +17,7 @@ service UserDefinedFunction {
   rpc MapTFn(Datum) returns (DatumList);
 
   // ReduceFn applies a reduce function to a datum stream.
-  rpc ReduceFn(stream Datum) returns (DatumList);
+  rpc ReduceFn(stream Datum) returns (stream DatumList);
 
   // IsReady is the heartbeat endpoint for gRPC.
   rpc IsReady(google.protobuf.Empty) returns (ReadyResponse);

src/test/java/io/numaproj/numaflow/function/FunctionServerTest.java

@@ -165,7 +165,7 @@ public void reducerWithOneKey() {
         String expectedKey = reduceKey + REDUCE_PROCESSED_KEY_SUFFIX;
         // sum of first 10 numbers 1 to 10 -> 55
         ByteString expectedValue = ByteString.copyFromUtf8(String.valueOf(55));
-        while (outputStreamObserver.resultDatum.get() == null) ;
+        while (!outputStreamObserver.completed.get());
 
         assertEquals(1, outputStreamObserver.resultDatum.get().getElementsCount());
         assertEquals(expectedKey, outputStreamObserver.resultDatum.get().getElements(0).getKey());
@@ -211,8 +211,7 @@ public void reducerWithMultipleKey() {
         // sum of first 10 numbers 1 to 10 -> 55
         ByteString expectedValue = ByteString.copyFromUtf8(String.valueOf(55));
 
-        while (outputStreamObserver.resultDatum.get() == null) ;
-
+        while (!outputStreamObserver.completed.get());
         Udfunction.DatumList result = outputStreamObserver.resultDatum.get();
         assertEquals(100, result.getElementsCount());
         for (int i = 0; i < keyCount; i++) {

src/test/java/io/numaproj/numaflow/function/ReduceOutputStreamObserver.java

@@ -8,12 +8,13 @@
 
 @Slf4j
 public class ReduceOutputStreamObserver implements StreamObserver<Udfunction.DatumList> {
-    public AtomicReference<Udfunction.DatumList> resultDatum = new AtomicReference<>();
+    public AtomicReference<Boolean> completed = new AtomicReference<>(false);
+    public AtomicReference<Udfunction.DatumList> resultDatum = new AtomicReference<>(Udfunction.DatumList.newBuilder().build());
     public Throwable t;
 
     @Override
     public void onNext(Udfunction.DatumList datum) {
-        resultDatum.set(datum);
+        resultDatum.set(resultDatum.get().toBuilder().addAllElements(datum.getElementsList()).build());
     }
 
     @Override
@@ -24,5 +25,6 @@ public void onError(Throwable throwable) {
     @Override
     public void onCompleted() {
         log.info("on completed executed");
+        this.completed.set(true);
     }
 }