io.numaproj.numaflow.reducestreamer.Service Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of numaflow-java Show documentation
SDK to implement Numaflow Source or User Defined Functions or Sinks in Java.
The newest version!
package io.numaproj.numaflow.reducestreamer;

import akka.actor.ActorRef;
import akka.actor.ActorSystem;
import akka.actor.AllDeadLetters;
import com.google.protobuf.Empty;
import io.grpc.Status;
import io.grpc.stub.StreamObserver;
import io.numaproj.numaflow.reduce.v1.ReduceGrpc;
import io.numaproj.numaflow.reduce.v1.ReduceOuterClass;
import io.numaproj.numaflow.reducestreamer.model.IntervalWindow;
import io.numaproj.numaflow.reducestreamer.model.Metadata;
import io.numaproj.numaflow.reducestreamer.model.ReduceStreamer;
import io.numaproj.numaflow.reducestreamer.model.ReduceStreamerFactory;
import io.numaproj.numaflow.shared.GrpcServerUtils;
import lombok.extern.slf4j.Slf4j;

import java.time.Instant;
import java.util.concurrent.CompletableFuture;

import static io.numaproj.numaflow.reduce.v1.ReduceGrpc.getReduceFnMethod;

@Slf4j
class Service extends ReduceGrpc.ReduceImplBase {
    public static final ActorSystem reduceActorSystem = ActorSystem.create("reducestream");

    private final ReduceStreamerFactory reduceStreamerFactory;

    public Service(ReduceStreamerFactory reduceStreamerFactory) {
        this.reduceStreamerFactory = reduceStreamerFactory;
    }

    static void handleFailure(
            CompletableFuture failureFuture,
            StreamObserver responseObserver) {
        new Thread(() -> {
            try {
                failureFuture.get();
            } catch (Exception e) {
                e.printStackTrace();
                var status = Status.UNKNOWN.withDescription(e.getMessage()).withCause(e);
                responseObserver.onError(status.asException());
            }
        }).start();
    }

    /**
     * Streams input data to reduceFn and returns the result.
     */
    @Override
    public StreamObserver reduceFn(final StreamObserver responseObserver) {
        if (this.reduceStreamerFactory == null) {
            return io.grpc.stub.ServerCalls.asyncUnimplementedStreamingCall(
                    getReduceFnMethod(),
                    responseObserver);
        }

        // get window start and end time from gPRC metadata
        String winSt = GrpcServerUtils.WINDOW_START_TIME.get();
        String winEt = GrpcServerUtils.WINDOW_END_TIME.get();

        // convert the start and end time to Instant
        Instant startTime = Instant.ofEpochMilli(Long.parseLong(winSt));
        Instant endTime = Instant.ofEpochMilli(Long.parseLong(winEt));

        // create metadata
        IntervalWindow iw = new IntervalWindowImpl(startTime, endTime);
        Metadata md = new MetadataImpl(iw);

        CompletableFuture failureFuture = new CompletableFuture<>();

        // create a shutdown actor that listens to exceptions.
        ActorRef shutdownActorRef = reduceActorSystem.
                actorOf(ShutdownActor.props(failureFuture));

        // subscribe for dead letters
        reduceActorSystem.getEventStream().subscribe(shutdownActorRef, AllDeadLetters.class);

        handleFailure(failureFuture, responseObserver);

        // create an output actor that ensures synchronized delivery of reduce responses.
        ActorRef outputActor = reduceActorSystem.
                actorOf(OutputActor.props(responseObserver));
        /*
            create a supervisor actor which assign the tasks to child actors.
            we create a child actor for every unique set of keys in a window.
        */
        ActorRef supervisorActor = reduceActorSystem
                .actorOf(SupervisorActor.props(
                        reduceStreamerFactory,
                        md,
                        shutdownActorRef,
                        outputActor));


        return new StreamObserver<>() {
            @Override
            public void onNext(ReduceOuterClass.ReduceRequest datum) {
                // send the message to parent actor, which takes care of distribution.
                if (!supervisorActor.isTerminated()) {
                    supervisorActor.tell(new ActorRequest(datum), ActorRef.noSender());
                } else {
                    responseObserver.onError(new Throwable("Supervisor actor was terminated"));
                }
            }

            @Override
            public void onError(Throwable throwable) {
                log.error("Error from the client - {}", throwable.getMessage());
                responseObserver.onError(throwable);
            }

            @Override
            public void onCompleted() {
                // indicate the end of input to the supervisor
                supervisorActor.tell(Constants.EOF, ActorRef.noSender());
            }
        };
    }

    /**
     * IsReady is the heartbeat endpoint for gRPC.
     */
    @Override
    public void isReady(
            Empty request,
            StreamObserver responseObserver) {
        responseObserver.onNext(ReduceOuterClass.ReadyResponse.newBuilder().setReady(true).build());
        responseObserver.onCompleted();
    }
}