io.numaproj.numaflow.function.FunctionService Maven / Gradle / Ivy

Go to download
package io.numaproj.numaflow.function;

import akka.actor.ActorRef;
import akka.actor.ActorSystem;
import akka.actor.AllDeadLetters;
import com.google.protobuf.ByteString;
import com.google.protobuf.Empty;
import com.google.protobuf.Timestamp;
import io.grpc.stub.StreamObserver;
import io.numaproj.numaflow.function.handlers.MapHandler;
import io.numaproj.numaflow.function.handlers.MapTHandler;
import io.numaproj.numaflow.function.handlers.ReduceHandler;
import io.numaproj.numaflow.function.handlers.ReducerFactory;
import io.numaproj.numaflow.function.interfaces.IntervalWindow;
import io.numaproj.numaflow.function.interfaces.Metadata;
import io.numaproj.numaflow.function.metadata.IntervalWindowImpl;
import io.numaproj.numaflow.function.metadata.MetadataImpl;
import io.numaproj.numaflow.function.types.MessageList;
import io.numaproj.numaflow.function.types.MessageTList;
import io.numaproj.numaflow.function.v1.Udfunction;
import io.numaproj.numaflow.function.v1.Udfunction.EventTime;
import io.numaproj.numaflow.function.v1.UserDefinedFunctionGrpc;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;

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

import static io.numaproj.numaflow.function.FunctionConstants.EOF;
import static io.numaproj.numaflow.function.v1.UserDefinedFunctionGrpc.getMapFnMethod;
import static io.numaproj.numaflow.function.v1.UserDefinedFunctionGrpc.getReduceFnMethod;

@Slf4j
@NoArgsConstructor
class FunctionService extends UserDefinedFunctionGrpc.UserDefinedFunctionImplBase {

    public static final ActorSystem functionActorSystem = ActorSystem.create("reduce");

    private MapHandler mapHandler;
    private MapTHandler mapTHandler;
    private ReducerFactory reducerFactory;

    public void setMapHandler(MapHandler mapHandler) {
        this.mapHandler = mapHandler;
    }

    public void setMapTHandler(MapTHandler mapTHandler) {
        this.mapTHandler = mapTHandler;
    }

    public void setReduceHandler(ReducerFactory reducerFactory) {
        this.reducerFactory = reducerFactory;
    }

    /**
     * Applies a function to each datum element.
     */
    @Override
    public void mapFn(
            Udfunction.DatumRequest request,
            StreamObserver responseObserver) {
        if (this.mapHandler == null) {
            io.grpc.stub.ServerCalls.asyncUnimplementedUnaryCall(
                    getMapFnMethod(),
                    responseObserver);
            return;
        }

        // get Datum from request
        HandlerDatumMetadata handlerDatumMetadata = new HandlerDatumMetadata(
                request.getMetadata().getId(),
                request.getMetadata().getNumDelivered()
        );
        HandlerDatum handlerDatum = new HandlerDatum(
                request.getValue().toByteArray(),
                Instant.ofEpochSecond(
                        request.getWatermark().getWatermark().getSeconds(),
                        request.getWatermark().getWatermark().getNanos()),
                Instant.ofEpochSecond(
                        request.getEventTime().getEventTime().getSeconds(),
                        request.getEventTime().getEventTime().getNanos()),
                handlerDatumMetadata
        );

        // process Datum
        MessageList messageList = mapHandler.processMessage(request
                .getKeysList()
                .toArray(new String[0]), handlerDatum);

        // set response
        responseObserver.onNext(buildDatumListResponse(messageList));
        responseObserver.onCompleted();
    }

    @Override
    public void mapTFn(
            Udfunction.DatumRequest request,
            StreamObserver responseObserver) {

        if (this.mapTHandler == null) {
            io.grpc.stub.ServerCalls.asyncUnimplementedUnaryCall(
                    getMapFnMethod(),
                    responseObserver);
            return;
        }

        // get Datum from request
        HandlerDatumMetadata handlerDatumMetadata = new HandlerDatumMetadata(
                request.getMetadata().getId(),
                request.getMetadata().getNumDelivered()
        );
        HandlerDatum handlerDatum = new HandlerDatum(
                request.getValue().toByteArray(),
                Instant.ofEpochSecond(
                        request.getWatermark().getWatermark().getSeconds(),
                        request.getWatermark().getWatermark().getNanos()),
                Instant.ofEpochSecond(
                        request.getEventTime().getEventTime().getSeconds(),
                        request.getEventTime().getEventTime().getNanos()),
                handlerDatumMetadata
        );

        // process Datum
        MessageTList messageTList = mapTHandler.processMessage(request
                .getKeysList()
                .toArray(new String[0]), handlerDatum);

        // set response
        responseObserver.onNext(buildDatumListResponse(messageTList));
        responseObserver.onCompleted();
    }

    /**
     * Streams input data to reduceFn and returns the result.
     */
    @Override
    public StreamObserver reduceFn(final StreamObserver responseObserver) {

        if (this.reducerFactory == null) {
            return io.grpc.stub.ServerCalls.asyncUnimplementedStreamingCall(
                    getReduceFnMethod(),
                    responseObserver);
        }

        // get window start and end time from gPRC metadata
        String winSt = FunctionConstants.WINDOW_START_TIME.get();
        String winEt = FunctionConstants.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 = functionActorSystem.
                actorOf(ReduceShutdownActor.props(responseObserver, failureFuture));

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

        handleFailure(failureFuture);
        /*
            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 = functionActorSystem
                .actorOf(ReduceSupervisorActor.props(
                        reducerFactory,
                        md,
                        shutdownActorRef,
                        responseObserver));


        return new StreamObserver() {
            @Override
            public void onNext(Udfunction.DatumRequest datum) {
                // send the message to parent actor, which takes care of distribution.
                if (!supervisorActor.isTerminated()) {
                    supervisorActor.tell(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(EOF, ActorRef.noSender());

            }
        };
    }

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

    private Udfunction.DatumResponseList buildDatumListResponse(MessageList messageList) {
        Udfunction.DatumResponseList.Builder datumListBuilder = Udfunction.DatumResponseList.newBuilder();
        messageList.getMessages().forEach(message -> {
            datumListBuilder.addElements(Udfunction.DatumResponse.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 datumListBuilder.build();
    }

    private Udfunction.DatumResponseList buildDatumListResponse(MessageTList messageTList) {
        Udfunction.DatumResponseList.Builder datumListBuilder = Udfunction.DatumResponseList.newBuilder();
        messageTList.getMessages().forEach(messageT -> {
            datumListBuilder.addElements(Udfunction.DatumResponse.newBuilder()
                    .setEventTime(
                            messageT.getEventTime() == null ? EventTime.newBuilder().setEventTime(
                                    Timestamp.getDefaultInstance()) : EventTime
                                    .newBuilder()
                                    .setEventTime
                                            (Timestamp.newBuilder()
                                                    .setSeconds(messageT
                                                            .getEventTime()
                                                            .getEpochSecond())
                                                    .setNanos(messageT.getEventTime().getNano()))
                    )
                    .addAllKeys(messageT.getKeys()
                            == null ? new ArrayList<>() : List.of(messageT.getKeys()))
                    .addAllTags(messageT.getTags()
                            == null ? new ArrayList<>() : List.of(messageT.getTags()))
                    .setValue(messageT.getValue() == null ? ByteString.EMPTY : ByteString.copyFrom(
                            messageT.getValue()))
                    .build());
        });
        return datumListBuilder.build();
    }

    // log the exception and exit if there are any uncaught exceptions.
    private void handleFailure(CompletableFuture failureFuture) {
        new Thread(() -> {
            try {
                failureFuture.get();
            } catch (Exception e) {
                e.printStackTrace();
                System.exit(1);
            }
        }).start();
    }
}