Skip to content

Commit

Permalink
feat: use gRPC Bidirectional Streaming for Mapper
Browse files Browse the repository at this point in the history
Signed-off-by: Yashash H L <[email protected]>
  • Loading branch information
yhl25 committed Oct 5, 2024
1 parent c387f9a commit 40ec7d3
Show file tree
Hide file tree
Showing 21 changed files with 568 additions and 154 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
2 changes: 2 additions & 0 deletions src/main/java/io/numaproj/numaflow/mapper/Constants.java
Original file line number Diff line number Diff line change
Expand Up @@ -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";
}
142 changes: 142 additions & 0 deletions src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java
Original file line number Diff line number Diff line change
@@ -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());
}

Check warning on line 76 in src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java#L69-L76

Added lines #L69 - L76 were not covered by tests

@Override
public void postStop() {
log.debug("post stop of supervisor executed - {}", getSelf().toString());
}

Check warning on line 81 in src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java#L80-L81

Added lines #L80 - L81 were not covered by tests

@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());
}

Check warning on line 124 in src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java#L120-L124

Added lines #L120 - L124 were not covered by tests

@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();

Check warning on line 137 in src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java#L132-L137

Added lines #L132 - L137 were not covered by tests
})
.build()
);
}
}
87 changes: 87 additions & 0 deletions src/main/java/io/numaproj/numaflow/mapper/MapperActor.java
Original file line number Diff line number Diff line change
@@ -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();
}
}
Loading

0 comments on commit 40ec7d3

Please sign in to comment.