diff --git a/examples/pom.xml b/examples/pom.xml
index 55a0020..a7e7b35 100644
--- a/examples/pom.xml
+++ b/examples/pom.xml
@@ -55,6 +55,9 @@
dockerBuild
+
+ amazoncorretto:11
+
io.numaproj.numaflow.examples.batchmap.flatmap.BatchFlatMap
@@ -74,6 +77,9 @@
dockerBuild
+
+ amazoncorretto:11
+
io.numaproj.numaflow.examples.sourcetransformer.eventtimefilter.EventTimeFilterFunction
@@ -93,6 +99,9 @@
dockerBuild
+
+ amazoncorretto:11
+
io.numaproj.numaflow.examples.mapstream.flatmapstream.FlatMapStreamFunction
@@ -110,6 +119,9 @@
dockerBuild
+
+ amazoncorretto:11
+
io.numaproj.numaflow.examples.map.flatmap.FlatMapFunction
@@ -195,6 +207,9 @@
dockerBuild
+
+ amazoncorretto:11
+
io.numaproj.numaflow.examples.map.forward.ForwardFunction
@@ -266,6 +281,9 @@
dockerBuild
+
+ amazoncorretto:11
+
io.numaproj.numaflow.examples.source.simple.SimpleSource
diff --git a/src/main/java/io/numaproj/numaflow/batchmapper/Server.java b/src/main/java/io/numaproj/numaflow/batchmapper/Server.java
index 83c276f..bca4d4f 100644
--- a/src/main/java/io/numaproj/numaflow/batchmapper/Server.java
+++ b/src/main/java/io/numaproj/numaflow/batchmapper/Server.java
@@ -6,10 +6,12 @@
import io.numaproj.numaflow.info.ContainerType;
import io.numaproj.numaflow.info.ServerInfoAccessor;
import io.numaproj.numaflow.info.ServerInfoAccessorImpl;
+import io.numaproj.numaflow.shared.GrpcServerHelper;
import io.numaproj.numaflow.shared.GrpcServerUtils;
import lombok.extern.slf4j.Slf4j;
import java.util.Collections;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
/**
@@ -20,8 +22,10 @@ public class Server {
private final GRPCConfig grpcConfig;
private final Service service;
+ private final CompletableFuture shutdownSignal;
private final ServerInfoAccessor serverInfoAccessor = new ServerInfoAccessorImpl(new ObjectMapper());
private io.grpc.Server server;
+ private final GrpcServerHelper grpcServerHelper;
/**
* constructor to create sink gRPC server.
@@ -39,8 +43,10 @@ public Server(BatchMapper batchMapper) {
* @param batchMapper to process the message
*/
public Server(BatchMapper batchMapper, GRPCConfig grpcConfig) {
- this.service = new Service(batchMapper);
+ this.shutdownSignal = new CompletableFuture<>();
+ this.service = new Service(batchMapper, this.shutdownSignal);
this.grpcConfig = grpcConfig;
+ this.grpcServerHelper = new GrpcServerHelper();
}
/**
@@ -57,35 +63,55 @@ public void start() throws Exception {
Collections.singletonMap(Constants.MAP_MODE_KEY, Constants.MAP_MODE));
if (this.server == null) {
- // create server builder
- ServerBuilder> serverBuilder = GrpcServerUtils.createServerBuilder(
+ this.server = grpcServerHelper.createServer(
grpcConfig.getSocketPath(),
grpcConfig.getMaxMessageSize(),
grpcConfig.isLocal(),
- grpcConfig.getPort());
- // build server
- this.server = serverBuilder
- .addService(this.service)
- .build();
+ grpcConfig.getPort(),
+ this.service);
}
- // start server
server.start();
log.info(
- "Server started, listening on socket path: " + grpcConfig.getSocketPath());
+ "server started, listening on socket path: " + grpcConfig.getSocketPath());
- // register shutdown hook
+ // register shutdown hook to gracefully shut down the server
Runtime.getRuntime().addShutdownHook(new Thread(() -> {
// Use stderr here since the logger may have been reset by its JVM shutdown hook.
System.err.println("*** shutting down gRPC server since JVM is shutting down");
+ if (server != null && server.isTerminated()) {
+ return;
+ }
try {
Server.this.stop();
+ log.info("gracefully shutting down event loop groups");
+ this.grpcServerHelper.gracefullyShutdownEventLoopGroups();
} catch (InterruptedException e) {
Thread.interrupted();
e.printStackTrace(System.err);
}
}));
+
+ // if there are any exceptions, shutdown the server gracefully.
+ shutdownSignal.whenCompleteAsync((v, e) -> {
+ if (server != null && server.isTerminated()) {
+ return;
+ }
+
+ if (e != null) {
+ System.err.println("*** shutting down batch map gRPC server because of an exception - " + e.getMessage());
+ try {
+ log.info("stopping server");
+ Server.this.stop();
+ log.info("gracefully shutting down event loop groups");
+ this.grpcServerHelper.gracefullyShutdownEventLoopGroups();
+ } catch (InterruptedException ex) {
+ Thread.interrupted();
+ ex.printStackTrace(System.err);
+ }
+ }
+ });
}
/**
@@ -96,7 +122,9 @@ public void start() throws Exception {
* @throws InterruptedException if the current thread is interrupted while waiting
*/
public void awaitTermination() throws InterruptedException {
+ log.info("batch map server is waiting for termination");
server.awaitTermination();
+ log.info("batch map server has terminated");
}
/**
diff --git a/src/main/java/io/numaproj/numaflow/batchmapper/Service.java b/src/main/java/io/numaproj/numaflow/batchmapper/Service.java
index b582a9b..fd967cd 100644
--- a/src/main/java/io/numaproj/numaflow/batchmapper/Service.java
+++ b/src/main/java/io/numaproj/numaflow/batchmapper/Service.java
@@ -35,6 +35,9 @@ class Service extends MapGrpc.MapImplBase {
// 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) {
@@ -93,8 +96,9 @@ public void onNext(MapOuterClass.MapRequest mapRequest) {
datumStream.writeMessage(constructHandlerDatum(mapRequest));
}
} catch (Exception e) {
- log.error("Encountered an error in batch map", e);
- responseObserver.onError(Status.UNKNOWN
+ log.error("Encountered an error in batch map onNext - {}", e.getMessage());
+ shutdownSignal.completeExceptionally(e);
+ responseObserver.onError(Status.INTERNAL
.withDescription(e.getMessage())
.withCause(e)
.asException());
@@ -104,11 +108,12 @@ public void onNext(MapOuterClass.MapRequest mapRequest) {
// Called when an error occurs
@Override
public void onError(Throwable throwable) {
- log.error("Error Encountered in batchMap Stream", throwable);
- var status = Status.UNKNOWN
+ log.error("Error Encountered in batchMap Stream - {}", throwable.getMessage());
+ shutdownSignal.completeExceptionally(throwable);
+ responseObserver.onError(Status.INTERNAL
.withDescription(throwable.getMessage())
- .withCause(throwable);
- responseObserver.onError(status.asException());
+ .withCause(throwable)
+ .asException());
}
// Called when the client has finished sending requests
diff --git a/src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java b/src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java
index 6a1649e..3ce663b 100644
--- a/src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java
+++ b/src/main/java/io/numaproj/numaflow/mapper/MapSupervisorActor.java
@@ -46,7 +46,9 @@
class MapSupervisorActor extends AbstractActor {
private final Mapper mapper;
private final StreamObserver responseObserver;
- private final CompletableFuture failureFuture;
+ private final CompletableFuture shutdownSignal;
+ private int activeMapperCount;
+ private Exception userException;
public MapSupervisorActor(
Mapper mapper,
@@ -54,25 +56,28 @@ public MapSupervisorActor(
CompletableFuture failureFuture) {
this.mapper = mapper;
this.responseObserver = responseObserver;
- this.failureFuture = failureFuture;
+ this.shutdownSignal = failureFuture;
+ this.userException = null;
+ this.activeMapperCount = 0;
}
public static Props props(
Mapper mapper,
StreamObserver responseObserver,
- CompletableFuture failureFuture) {
- return Props.create(MapSupervisorActor.class, mapper, responseObserver, failureFuture);
+ CompletableFuture shutdownSignal) {
+ return Props.create(MapSupervisorActor.class, mapper, responseObserver, shutdownSignal);
}
@Override
public void preRestart(Throwable reason, Optional