use of org.apache.beam.sdk.fn.data.BeamFnDataOutboundAggregator in project beam by apache.
the class PTransformRunnerFactoryTestContext method addOutgoingDataEndpoint.
@Override
public <T> FnDataReceiver<T> addOutgoingDataEndpoint(ApiServiceDescriptor apiServiceDescriptor, Coder<T> coder) {
BeamFnDataOutboundAggregator aggregator = getOutboundAggregators().get(apiServiceDescriptor);
FnDataReceiver<T> receiver = aggregator.registerOutputDataLocation(getPTransformId(), coder);
getOutgoingDataEndpoints().computeIfAbsent(apiServiceDescriptor, (unused) -> new ArrayList<>()).add(DataEndpoint.create(getPTransformId(), coder, receiver));
return receiver;
}
use of org.apache.beam.sdk.fn.data.BeamFnDataOutboundAggregator in project beam by apache.
the class BeamFnDataGrpcClientTest method testForOutboundConsumer.
@Test
public void testForOutboundConsumer() throws Exception {
CountDownLatch waitForInboundServerValuesCompletion = new CountDownLatch(2);
Collection<BeamFnApi.Elements> inboundServerValues = new ConcurrentLinkedQueue<>();
CallStreamObserver<BeamFnApi.Elements> inboundServerObserver = TestStreams.withOnNext((BeamFnApi.Elements t) -> {
inboundServerValues.add(t);
waitForInboundServerValuesCompletion.countDown();
}).build();
Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID()).build();
Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnDataGrpc.BeamFnDataImplBase() {
@Override
public StreamObserver<BeamFnApi.Elements> data(StreamObserver<BeamFnApi.Elements> outboundObserver) {
return inboundServerObserver;
}
}).build();
server.start();
try {
ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
BeamFnDataGrpcClient clientFactory = new BeamFnDataGrpcClient(PipelineOptionsFactory.fromArgs(new String[] { "--experiments=data_buffer_size_limit=20" }).create(), (Endpoints.ApiServiceDescriptor descriptor) -> channel, OutboundObserverFactory.trivial());
BeamFnDataOutboundAggregator aggregator = clientFactory.createOutboundAggregator(apiServiceDescriptor, () -> INSTRUCTION_ID_A, false);
FnDataReceiver<WindowedValue<String>> fnDataReceiver = aggregator.registerOutputDataLocation(TRANSFORM_ID_A, CODER);
fnDataReceiver.accept(valueInGlobalWindow("ABC"));
fnDataReceiver.accept(valueInGlobalWindow("DEF"));
fnDataReceiver.accept(valueInGlobalWindow("GHI"));
aggregator.sendOrCollectBufferedDataAndFinishOutboundStreams();
waitForInboundServerValuesCompletion.await();
assertThat(inboundServerValues, contains(ELEMENTS_A_1, ELEMENTS_A_2));
} finally {
server.shutdownNow();
}
}
Aggregations