
io.numaproj.numaflow.batchmapper.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
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.batchmapper;
import com.google.protobuf.ByteString;
import com.google.protobuf.Empty;
import io.grpc.Status;
import io.grpc.stub.StreamObserver;
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 java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
@Slf4j
@AllArgsConstructor
class Service extends MapGrpc.MapImplBase {
// Executor service for the batch mapper. It is a fixed size thread pool
// with the number of threads equal to the number of cores on the machine times 2.
// We use 2 times the number of cores because the batch mapper is a CPU intensive task.
private final ExecutorService mapTaskExecutor = Executors
.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2);
// Time to wait for the batch mapper to shut down, in seconds.
// We use 30 seconds as the default value because it provides a balance between giving tasks enough time to complete
// and not delaying program termination unduly.
private final long SHUTDOWN_TIME = 30;
// BatchMapper instance to process the messages
private final BatchMapper batchMapper;
// Signal to shut down the gRPC server
private final CompletableFuture shutdownSignal;
// Applies a map function to each datum element in the stream.
@Override
public StreamObserver mapFn(StreamObserver responseObserver) {
// If the batchMapper is null, return an unimplemented call
if (this.batchMapper == null) {
return io.grpc.stub.ServerCalls.asyncUnimplementedStreamingCall(
MapGrpc.getMapFnMethod(),
responseObserver);
}
// Return a new StreamObserver to handle the incoming requests
return new StreamObserver<>() {
private boolean startOfStream = true;
private boolean handshakeDone = false;
private DatumIteratorImpl datumStream;
private CompletableFuture result;
// Called for each incoming request
@Override
public void onNext(MapOuterClass.MapRequest mapRequest) {
try {
// Make sure the handshake is done before processing the messages
if (!handshakeDone) {
if (!mapRequest.hasHandshake() || !mapRequest.getHandshake().getSot()) {
responseObserver.onError(Status.INVALID_ARGUMENT
.withDescription("Handshake request not received")
.asException());
return;
}
responseObserver.onNext(MapOuterClass.MapResponse.newBuilder()
.setHandshake(mapRequest.getHandshake())
.build());
handshakeDone = true;
return;
}
// Create a DatumIterator to write the messages to the batch mapper
// and start the batch mapper if it is the start of the stream
if (startOfStream) {
datumStream = new DatumIteratorImpl();
result = CompletableFuture.supplyAsync(
() -> batchMapper.processMessage(datumStream),
mapTaskExecutor);
startOfStream = false;
}
// If end of transmission, write EOF datum to the stream
// Wait for the result and send the response back to the client
if (mapRequest.hasStatus() && mapRequest.getStatus().getEot()) {
datumStream.writeMessage(HandlerDatum.EOF_DATUM);
BatchResponses responses = result.join();
buildAndStreamResponse(responses, responseObserver);
startOfStream = true;
} else {
datumStream.writeMessage(constructHandlerDatum(mapRequest));
}
} catch (Exception e) {
log.error("Encountered an error in batch map onNext - {}", e.getMessage());
shutdownSignal.completeExceptionally(e);
responseObserver.onError(Status.INTERNAL
.withDescription(e.getMessage())
.withCause(e)
.asException());
}
}
// Called when an error occurs
@Override
public void onError(Throwable throwable) {
log.error("Error Encountered in batchMap Stream - {}", throwable.getMessage());
shutdownSignal.completeExceptionally(throwable);
responseObserver.onError(Status.INTERNAL
.withDescription(throwable.getMessage())
.withCause(throwable)
.asException());
}
// Called when the client has finished sending requests
@Override
public void onCompleted() {
responseObserver.onCompleted();
}
};
}
// Build and stream the response back to the client
private void buildAndStreamResponse(
BatchResponses responses,
StreamObserver responseObserver) {
responses.getItems().forEach(message -> {
List mapResponseResult = new ArrayList<>();
message.getItems().forEach(res -> {
mapResponseResult.add(
MapOuterClass.MapResponse.Result
.newBuilder()
.setValue(res.getValue()
== null ? ByteString.EMPTY : ByteString.copyFrom(
res.getValue()))
.addAllKeys(res.getKeys()
== null ? new ArrayList<>() : List.of(res.getKeys()))
.addAllTags(res.getTags()
== null ? new ArrayList<>() : List.of(res.getTags()))
.build()
);
});
MapOuterClass.MapResponse singleRequestResponse = MapOuterClass.MapResponse
.newBuilder()
.setId(message.getId())
.addAllResults(mapResponseResult)
.build();
responseObserver.onNext(singleRequestResponse);
});
// 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);
responseObserver.onCompleted();
}
// IsReady is the heartbeat endpoint for gRPC.
@Override
public void isReady(
Empty request,
StreamObserver responseObserver) {
responseObserver.onNext(MapOuterClass.ReadyResponse.newBuilder().setReady(true).build());
responseObserver.onCompleted();
}
// Construct a HandlerDatum from a MapRequest
private HandlerDatum constructHandlerDatum(MapOuterClass.MapRequest d) {
return new HandlerDatum(
d.getRequest().getKeysList().toArray(new String[0]),
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.getId(),
d.getRequest().getHeadersMap()
);
}
// Shuts down the executor service
public void shutDown() {
this.mapTaskExecutor.shutdown();
try {
if (!mapTaskExecutor.awaitTermination(SHUTDOWN_TIME, TimeUnit.SECONDS)) {
log.error("BatchMap executor did not terminate in the specified time.");
List droppedTasks = mapTaskExecutor.shutdownNow();
log.error(
"BatchMap executor was abruptly shut down. {} tasks will not be executed.",
droppedTasks.size());
} else {
log.info("BatchMap executor was terminated.");
}
} catch (InterruptedException e) {
Thread.interrupted();
e.printStackTrace();
}
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy