-
Notifications
You must be signed in to change notification settings - Fork 10
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: Unify MapStream and Unary Map Operations Using a Shared gRPC Pr…
…otocol (#146) Signed-off-by: Yashash H L <[email protected]>
- Loading branch information
Showing
15 changed files
with
231 additions
and
137 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
117 changes: 79 additions & 38 deletions
117
src/main/java/io/numaproj/numaflow/mapstreamer/Service.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,64 +1,105 @@ | ||
| package io.numaproj.numaflow.mapstreamer; | ||
|
|
||
| import com.google.protobuf.Empty; | ||
| import io.grpc.Status; | ||
| import io.grpc.stub.StreamObserver; | ||
| import io.numaproj.numaflow.mapstream.v1.MapStreamGrpc; | ||
| import io.numaproj.numaflow.mapstream.v1.Mapstream; | ||
| import io.numaproj.numaflow.map.v1.MapGrpc; | ||
| import io.numaproj.numaflow.map.v1.MapOuterClass; | ||
| import lombok.AllArgsConstructor; | ||
| import lombok.extern.slf4j.Slf4j; | ||
|
|
||
| import java.time.Instant; | ||
|
|
||
| import static io.numaproj.numaflow.mapstream.v1.MapStreamGrpc.getMapStreamFnMethod; | ||
|
|
||
| import static io.numaproj.numaflow.map.v1.MapGrpc.getMapFnMethod; | ||
|
|
||
| @Slf4j | ||
| @AllArgsConstructor | ||
| class Service extends MapStreamGrpc.MapStreamImplBase { | ||
| class Service extends MapGrpc.MapImplBase { | ||
|
|
||
| private final MapStreamer mapStreamer; | ||
|
|
||
| /** | ||
| * Applies a map stream function to each request. | ||
| */ | ||
| @Override | ||
| public void mapStreamFn( | ||
| Mapstream.MapStreamRequest request, | ||
| StreamObserver<Mapstream.MapStreamResponse> responseObserver) { | ||
| public StreamObserver<MapOuterClass.MapRequest> mapFn(StreamObserver<MapOuterClass.MapResponse> responseObserver) { | ||
|
|
||
| if (this.mapStreamer == null) { | ||
| io.grpc.stub.ServerCalls.asyncUnimplementedUnaryCall( | ||
| getMapStreamFnMethod(), | ||
| return io.grpc.stub.ServerCalls.asyncUnimplementedStreamingCall( | ||
| getMapFnMethod(), | ||
| responseObserver); | ||
| return; | ||
| } | ||
|
|
||
| HandlerDatum handlerDatum = new HandlerDatum( | ||
| request.getValue().toByteArray(), | ||
| Instant.ofEpochSecond( | ||
| request.getWatermark().getSeconds(), | ||
| request.getWatermark().getNanos()), | ||
| Instant.ofEpochSecond( | ||
| request.getEventTime().getSeconds(), | ||
| request.getEventTime().getNanos()), | ||
| request.getHeadersMap() | ||
| ); | ||
| return new StreamObserver<>() { | ||
| private boolean handshakeDone = false; | ||
|
|
||
| // process Datum | ||
| this.mapStreamer.processMessage(request | ||
| .getKeysList() | ||
| .toArray(new String[0]), handlerDatum, new OutputObserverImpl(responseObserver)); | ||
| @Override | ||
| public void onNext(MapOuterClass.MapRequest request) { | ||
| // make sure the handshake is done before processing the messages | ||
| if (!handshakeDone) { | ||
| if (!request.hasHandshake() || !request.getHandshake().getSot()) { | ||
| responseObserver.onError(Status.INVALID_ARGUMENT | ||
| .withDescription("Handshake request not received") | ||
| .asException()); | ||
| return; | ||
| } | ||
| responseObserver.onNext(MapOuterClass.MapResponse.newBuilder() | ||
| .setHandshake(request.getHandshake()) | ||
| .build()); | ||
| handshakeDone = true; | ||
| return; | ||
| } | ||
|
|
||
| responseObserver.onCompleted(); | ||
| } | ||
| try { | ||
| // process the message | ||
| mapStreamer.processMessage( | ||
| request | ||
| .getRequest() | ||
| .getKeysList() | ||
| .toArray(new String[0]), | ||
| constructHandlerDatum(request), | ||
| new OutputObserverImpl(responseObserver)); | ||
| } catch (Exception e) { | ||
| log.error("Error processing message", e); | ||
| responseObserver.onError(Status.UNKNOWN | ||
| .withDescription(e.getMessage()) | ||
| .asException()); | ||
| return; | ||
| } | ||
|
|
||
| /** | ||
| * IsReady is the heartbeat endpoint for gRPC. | ||
| */ | ||
| @Override | ||
| public void isReady(Empty request, StreamObserver<Mapstream.ReadyResponse> responseObserver) { | ||
| responseObserver.onNext(Mapstream.ReadyResponse.newBuilder().setReady(true).build()); | ||
| responseObserver.onCompleted(); | ||
| // Send an EOT message to indicate the end of the transmission for the batch. | ||
| MapOuterClass.MapResponse eotResponse = MapOuterClass.MapResponse | ||
| .newBuilder() | ||
| .setStatus(MapOuterClass.TransmissionStatus | ||
| .newBuilder() | ||
| .setEot(true) | ||
| .build()).build(); | ||
| responseObserver.onNext(eotResponse); | ||
| } | ||
|
|
||
| @Override | ||
| public void onError(Throwable throwable) { | ||
| log.error("Error Encountered in mapStream Stream", throwable); | ||
| var status = Status.UNKNOWN | ||
| .withDescription(throwable.getMessage()) | ||
| .withCause(throwable); | ||
| responseObserver.onError(status.asException()); | ||
| } | ||
|
|
||
| @Override | ||
| public void onCompleted() { | ||
| responseObserver.onCompleted(); | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| // Construct a HandlerDatum from a MapRequest | ||
| private HandlerDatum constructHandlerDatum(MapOuterClass.MapRequest d) { | ||
| return new HandlerDatum( | ||
| d.getRequest().getValue().toByteArray(), | ||
| Instant.ofEpochSecond( | ||
| d.getRequest().getWatermark().getSeconds(), | ||
| d.getRequest().getWatermark().getNanos()), | ||
| Instant.ofEpochSecond( | ||
| d.getRequest().getEventTime().getSeconds(), | ||
| d.getRequest().getEventTime().getNanos()), | ||
| d.getRequest().getHeadersMap() | ||
| ); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
41 changes: 41 additions & 0 deletions
41
src/test/java/io/numaproj/numaflow/mapstreamer/MapStreamOutputStreamObserver.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,41 @@ | ||
| package io.numaproj.numaflow.mapstreamer; | ||
|
|
||
| import io.grpc.stub.StreamObserver; | ||
| import io.numaproj.numaflow.map.v1.MapOuterClass; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.List; | ||
| import java.util.concurrent.CompletableFuture; | ||
|
|
||
| public class MapStreamOutputStreamObserver implements StreamObserver<MapOuterClass.MapResponse> { | ||
| List<MapOuterClass.MapResponse> mapResponses = new ArrayList<>(); | ||
| CompletableFuture<Void> done = new CompletableFuture<>(); | ||
| Integer responseCount; | ||
|
|
||
| public MapStreamOutputStreamObserver(Integer responseCount) { | ||
| this.responseCount = responseCount; | ||
| } | ||
|
|
||
| @Override | ||
| public void onNext(MapOuterClass.MapResponse mapResponse) { | ||
| System.out.println("Received response: " + mapResponse); | ||
| mapResponses.add(mapResponse); | ||
| if (mapResponses.size() == responseCount) { | ||
| done.complete(null); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public void onError(Throwable throwable) { | ||
| done.completeExceptionally(throwable); | ||
| } | ||
|
|
||
| @Override | ||
| public void onCompleted() { | ||
| done.complete(null); | ||
| } | ||
|
|
||
| public List<MapOuterClass.MapResponse> getMapResponses() { | ||
| return mapResponses; | ||
| } | ||
| } |
Oops, something went wrong.