Add a flag so that app can close when stream ends

If user sets grpc.exitOnComplete=true then the app will close
down after the end of the output flux (which is the same as the
end of the input flux for a Consumer).

See projectriff/riff#544 for a change
in the sidecar that supports this, and makes it possible to
declare a k8s Job with a sidecar and a function invoker.

Author: Dave Syer

src/main/java/io/projectriff/invoker/ApplicationRunner.java

@@ -152,6 +152,15 @@ public Object evaluate(String expression, Object root, Object... attrs) {
 		return parsed.getValue(context);
 	}
 
+	public boolean isRunning() {
+		if (this.app == null) {
+			return false;
+		}
+		Expression parsed = new SpelExpressionParser()
+				.parseExpression("context.isRunning()");
+		return parsed.getValue(this.app, Boolean.class);
+	}
+
 	@PreDestroy
 	public void close() {
 		closeContext();

src/main/java/io/projectriff/invoker/GrpcConfiguration.java

@@ -38,6 +38,8 @@
 import org.springframework.cloud.function.context.catalog.FunctionInspector;
 import org.springframework.cloud.function.core.FluxConsumer;
 import org.springframework.cloud.function.core.FluxSupplier;
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.event.ContextRefreshedEvent;
 import org.springframework.context.event.EventListener;
@@ -56,6 +58,7 @@ public class GrpcConfiguration {
 	private static final Log logger = LogFactory.getLog(GrpcConfiguration.class);
 	private Server server;
 	private int port = 10382;
+	private boolean exitOnComplete;
 
 	@Autowired
 	private Gson mapper;
@@ -65,6 +68,8 @@ public class GrpcConfiguration {
 	private FunctionInspector inspector;
 	@Autowired
 	private FunctionCatalog catalog;
+	@Autowired
+	private ConfigurableApplicationContext context;
 
 	public int getPort() {
 		return port;
@@ -75,6 +80,14 @@ public void setPort(int port) {
 		this.port = port;
 	}
 
+	public boolean isExitOnComplete() {
+		return exitOnComplete;
+	}
+
+	public void setExitOnComplete(boolean exitOnComplete) {
+		this.exitOnComplete = exitOnComplete;
+	}
+
 	/** Start serving requests. */
 	@EventListener(ContextRefreshedEvent.class)
 	public void start() {
@@ -96,7 +109,8 @@ public void start() {
 									function(function,
 											inspector.getRegistration(function)
 													.getType()),
-									this.mapper, inspector.getInputType(function),
+									this::maybeClose, this.mapper,
+									inspector.getInputType(function),
 									inspector.getOutputType(function),
 									inspector.isMessage(function)))
 					.build();
@@ -176,6 +190,19 @@ private Function<Flux<?>, Flux<?>> function(Object result, FunctionType type) {
 		return output;
 	}
 
+	private void maybeClose() {
+		if (this.exitOnComplete) {
+			ApplicationContext context = this.context;
+			while (context instanceof ConfigurableApplicationContext) {
+				ConfigurableApplicationContext closeable = (ConfigurableApplicationContext) context;
+				if (closeable.isRunning()) {
+					closeable.close();
+				}
+				context = context.getParent();
+			}
+		}
+	}
+
 	private static final class ConsumerAdapter
 			implements Function<Flux<Object>, Flux<Object>> {
 
@@ -204,9 +231,10 @@ public SupplierAdapter(Supplier<Publisher<?>> result) {
 			this.result = result;
 		}
 
+		@SuppressWarnings("unchecked")
 		@Override
 		public Flux<Object> apply(Flux<Object> input) {
-			return Flux.from(result.get());
+			return (Flux<Object>) Flux.from(result.get());
 		}
 
 	}

src/main/java/io/projectriff/invoker/JavaFunctionInvokerServer.java

@@ -50,11 +50,14 @@ public class JavaFunctionInvokerServer
 	private final Class<?> inputType;
 	private final Class<?> outputType;
 	private final Gson mapper;
+	private Runnable callback;
 
 	private static final Log logger = LogFactory.getLog(JavaFunctionInvokerServer.class);
 
-	public JavaFunctionInvokerServer(Function<Flux<?>, Flux<?>> function, Gson mapper,
-			Class<?> inputType, Class<?> outputType, boolean isMessage) {
+	public JavaFunctionInvokerServer(Function<Flux<?>, Flux<?>> function,
+			Runnable callback, Gson mapper, Class<?> inputType, Class<?> outputType,
+			boolean isMessage) {
+		this.callback = callback;
 		this.mapper = mapper;
 		this.inputType = inputType;
 		this.outputType = outputType;
@@ -122,8 +125,11 @@ public StreamObserver<io.projectriff.grpc.function.FunctionProtos.Message> call(
 				}, () -> {
 					// Make sure the emitter is disposed (should work even if it already
 					// was since it's idempotent)
-					emitter.dispose();
+					if (!emitter.isDisposed()) {
+						emitter.dispose();
+					}
 					responseObserver.onCompleted();
+					this.callback.run();
 				});
 
 		return new StreamObserver<io.projectriff.grpc.function.FunctionProtos.Message>() {

src/test/java/io/projectriff/functions/Logger.java

@@ -27,7 +27,7 @@ public class Logger implements Consumer<Flux<String>> {
 
 	@Override
 	public void accept(Flux<String> input) {
-		input.map(value -> "Hello " + value).subscribe(System.out::println);
+		input.map(value -> "Hello " + value).doOnNext(System.out::println);
 	}
 
 }

src/test/java/io/projectriff/invoker/GrpcSinkTests.java

@@ -65,5 +65,18 @@ public void fluxConsumer() throws Exception {
 		ApplicationRunner runner = (ApplicationRunner) ReflectionTestUtils
 				.getField(this.runner, "runner");
 		assertThat(runner.containsBean("io.projectriff.functions.Logger")).isFalse();
+		assertThat(runner.isRunning()).isTrue();
+	}
+
+	@Test
+	public void fluxConsumerCloses() throws Exception {
+		runner.run("--server.port=0", "--grpc.port=" + port, "--grpc.exitOnComplete=true",
+				"--function.uri=file:target/test-classes"
+						+ "?handler=io.projectriff.functions.Logger");
+		List<String> result = client.send("foo");
+		assertThat(result).isEmpty();
+		ApplicationRunner runner = (ApplicationRunner) ReflectionTestUtils
+				.getField(this.runner, "runner");
+		assertThat(runner.isRunning()).isFalse();
 	}
 }

src/test/java/io/projectriff/invoker/GrpcSourceTests.java

@@ -65,5 +65,18 @@ public void fluxSupplier() throws Exception {
 		ApplicationRunner runner = (ApplicationRunner) ReflectionTestUtils
 				.getField(this.runner, "runner");
 		assertThat(runner.containsBean("io.projectriff.functions.Words")).isFalse();
+		assertThat(runner.isRunning()).isTrue();
+	}
+
+	@Test
+	public void fluxSupplierCloses() throws Exception {
+		runner.run("--server.port=0", "--grpc.port=" + port, "--grpc.exitOnComplete=true",
+				"--function.uri=file:target/test-classes"
+						+ "?handler=io.projectriff.functions.Words");
+		List<String> result = client.send();
+		assertThat(result).contains("foo");
+		ApplicationRunner runner = (ApplicationRunner) ReflectionTestUtils
+				.getField(this.runner, "runner");
+		assertThat(runner.isRunning()).isFalse();
 	}
 }