GH-744 Add initial biStreaming support for Function<Flux, Flux> in gRPC module

This commit is contained in:
Oleg Zhurakousky
2021-09-21 13:40:09 +02:00
parent 28b1c86849
commit 70a789da69
5 changed files with 103 additions and 13 deletions

View File

@@ -20,6 +20,7 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.cloud.function.context.FunctionCatalog;
import org.springframework.cloud.function.context.FunctionProperties;
import org.springframework.cloud.function.grpc.MessagingServiceGrpc.MessagingServiceImplBase;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
@@ -31,16 +32,16 @@ import org.springframework.context.annotation.Configuration;
@Configuration(proxyBeanMethods = false)
@EnableConfigurationProperties(FunctionGrpcProperties.class)
@ConditionalOnProperty(name = "spring.cloud.function.grpc.mode", havingValue = "server", matchIfMissing = false)
public class GrpcAutoConfiguration {
class GrpcAutoConfiguration {
@Bean
public GrpcServer grpcServer(FunctionGrpcProperties grpcProperties, GrpcMessagingServiceImpl grpcMessagingService) {
public GrpcServer grpcServer(FunctionGrpcProperties grpcProperties, MessagingServiceImplBase grpcMessagingService) {
return new GrpcServer(grpcProperties, grpcMessagingService);
}
@Bean
public GrpcMessagingServiceImpl grpcMessageService(FunctionProperties funcProperties, FunctionCatalog functionCatalog) {
return new GrpcMessagingServiceImpl(funcProperties, functionCatalog);
public GrpcServerMessageHandler grpcMessageService(FunctionProperties funcProperties, FunctionCatalog functionCatalog) {
return new GrpcServerMessageHandler(funcProperties, functionCatalog);
}
}

View File

@@ -38,6 +38,10 @@ import java.util.concurrent.atomic.AtomicBoolean;
import io.grpc.Status;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;
import reactor.core.publisher.Sinks.Many;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
@@ -46,6 +50,7 @@ import org.springframework.cloud.function.context.FunctionProperties;
import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper;
import org.springframework.cloud.function.grpc.MessagingServiceGrpc.MessagingServiceImplBase;
import org.springframework.messaging.Message;
import org.springframework.messaging.support.MessageBuilder;
import org.springframework.util.Assert;
/**
@@ -54,13 +59,13 @@ import org.springframework.util.Assert;
* @since 3.2
*
*/
class GrpcMessagingServiceImpl extends MessagingServiceImplBase {
class GrpcServerMessageHandler extends MessagingServiceImplBase {
private Log logger = LogFactory.getLog(GrpcMessagingServiceImpl.class);
private Log logger = LogFactory.getLog(GrpcServerMessageHandler.class);
private final FunctionInvocationWrapper function;
GrpcMessagingServiceImpl(FunctionProperties funcProperties, FunctionCatalog functionCatalog) {
GrpcServerMessageHandler(FunctionProperties funcProperties, FunctionCatalog functionCatalog) {
this.function = functionCatalog.lookup(funcProperties.getDefinition(), "application/json");
Assert.notNull(this.function, "Failed to lookup function " + funcProperties.getDefinition());
}
@@ -91,7 +96,6 @@ class GrpcMessagingServiceImpl extends MessagingServiceImplBase {
// }
//
@Override
@SuppressWarnings("unchecked")
public StreamObserver<GrpcMessage> biStream(StreamObserver<GrpcMessage> responseObserver) {
ServerCallStreamObserver<GrpcMessage> serverCallStreamObserver = (ServerCallStreamObserver<GrpcMessage>) responseObserver;
serverCallStreamObserver.disableAutoInboundFlowControl();
@@ -104,15 +108,26 @@ class GrpcMessagingServiceImpl extends MessagingServiceImplBase {
serverCallStreamObserver.request(1);
}
});
if (function.isInputTypePublisher()) {
return this.biStreamReactive(responseObserver, serverCallStreamObserver);
}
else {
return this.biStreamImperative(responseObserver, serverCallStreamObserver, wasReady);
}
}
private StreamObserver<GrpcMessage> biStreamImperative(StreamObserver<GrpcMessage> responseObserver,
ServerCallStreamObserver<GrpcMessage> serverCallStreamObserver, AtomicBoolean wasReady) {
return new StreamObserver<GrpcMessage>() {
@SuppressWarnings("unchecked")
@Override
public void onNext(GrpcMessage request) {
try {
Message<byte[]> message = GrpcUtils.fromGrpcMessage(request);
Message<byte[]> replyMessage = (Message<byte[]>) function
.apply(message);
Message<byte[]> replyMessage = (Message<byte[]>) function.apply(message);
GrpcMessage reply = GrpcUtils.toGrpcMessage(replyMessage);
@@ -147,4 +162,40 @@ class GrpcMessagingServiceImpl extends MessagingServiceImplBase {
}
};
}
@SuppressWarnings("unchecked")
private StreamObserver<GrpcMessage> biStreamReactive(StreamObserver<GrpcMessage> responseObserver,
ServerCallStreamObserver<GrpcMessage> serverCallStreamObserver) {
Many<Message<byte[]>> sink = Sinks.many().unicast().onBackpressureBuffer();
Flux<Message<byte[]>> flux = sink.asFlux();
Flux<Message<byte[]>> connectedFlux = (Flux<Message<byte[]>>) function.apply(flux);
connectedFlux.subscribe(functionResult -> {
GrpcMessage reply = GrpcUtils.toGrpcMessage(functionResult);
responseObserver.onNext(reply);
});
return new StreamObserver<GrpcMessage>() {
@Override
public void onNext(GrpcMessage value) {
sink.tryEmitNext(GrpcUtils.fromGrpcMessage(value));
serverCallStreamObserver.request(1);
}
@Override
public void onError(Throwable t) {
t.printStackTrace();
responseObserver.onCompleted();
}
@Override
public void onCompleted() {
logger.info("Server stream is complete");
sink.tryEmitComplete();
responseObserver.onCompleted();
}
};
}
}

View File

@@ -39,7 +39,7 @@ import org.springframework.messaging.support.MessageBuilder;
* @since 3.2
*
*/
public final class GrpcUtils {
final class GrpcUtils {
private static Log logger = LogFactory.getLog(GrpcUtils.class);
@@ -155,6 +155,7 @@ public final class GrpcUtils {
@Override
public void onNext(GrpcMessage message) {
System.out.println("RECEIVED: " + message);
if (logger.isDebugEnabled()) {
logger.debug("Receiving message: " + message);
}

View File

@@ -64,7 +64,7 @@ public class GrpcInteractionTests {
}
@Test
public void testBidirectionalStream() {
public void testBidirectionalStreamWithImperativeFunction() {
try (ConfigurableApplicationContext context = new SpringApplicationBuilder(
SampleConfiguration.class).web(WebApplicationType.NONE).run(
"--spring.jmx.enabled=false",
@@ -95,6 +95,38 @@ public class GrpcInteractionTests {
}
}
@Test
public void testBidirectionalStreamWithReactiveFunction() {
try (ConfigurableApplicationContext context = new SpringApplicationBuilder(
SampleConfiguration.class).web(WebApplicationType.NONE).run(
"--spring.jmx.enabled=false",
"--spring.cloud.function.definition=uppercaseReactive",
"--spring.cloud.function.grpc.port="
+ FunctionGrpcProperties.GRPC_PORT,
"--spring.cloud.function.grpc.mode=server")) {
List<Message<byte[]>> messages = new ArrayList<>();
messages.add(MessageBuilder.withPayload("\"Ricky\"".getBytes()).setHeader("foo", "bar")
.setHeader(MessageHeaders.CONTENT_TYPE, MimeTypeUtils.TEXT_PLAIN)
.build());
messages.add(MessageBuilder.withPayload("\"Julien\"".getBytes()).setHeader("foo", "bar")
.setHeader(MessageHeaders.CONTENT_TYPE, MimeTypeUtils.TEXT_PLAIN)
.build());
messages.add(MessageBuilder.withPayload("\"Bubbles\"".getBytes()).setHeader("foo", "bar")
.setHeader(MessageHeaders.CONTENT_TYPE, MimeTypeUtils.TEXT_PLAIN)
.build());
Flux<Message<byte[]>> clientResponseObserver =
GrpcUtils.biStreaming("localhost", FunctionGrpcProperties.GRPC_PORT, Flux.fromIterable(messages));
List<Message<byte[]>> results = clientResponseObserver.collectList().block(Duration.ofSeconds(5));
assertThat(results.size()).isEqualTo(3);
assertThat(results.get(0).getPayload()).isEqualTo("\"RICKY\"".getBytes());
assertThat(results.get(1).getPayload()).isEqualTo("\"JULIEN\"".getBytes());
assertThat(results.get(2).getPayload()).isEqualTo("\"BUBBLES\"".getBytes());
}
}
@EnableAutoConfiguration
public static class SampleConfiguration {
@@ -102,5 +134,10 @@ public class GrpcInteractionTests {
public Function<String, String> uppercase() {
return v -> v.toUpperCase();
}
@Bean
public Function<Flux<String>, Flux<String>> uppercaseReactive() {
return flux -> flux.map(v -> v.toUpperCase());
}
}
}