use of org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements.Data in project beam by apache.
the class BeamFnDataGrpcServiceTest method testMessageReceivedBySingleClientWhenThereAreMultipleClients.
@Test
public void testMessageReceivedBySingleClientWhenThereAreMultipleClients() throws Exception {
BlockingQueue<Elements> clientInboundElements = new LinkedBlockingQueue<>();
ExecutorService executorService = Executors.newCachedThreadPool();
CountDownLatch waitForInboundElements = new CountDownLatch(1);
int numberOfClients = 3;
for (int client = 0; client < numberOfClients; ++client) {
executorService.submit(() -> {
ManagedChannel channel = ManagedChannelFactory.createDefault().withInterceptors(Arrays.asList(AddHarnessIdInterceptor.create(WORKER_ID))).forDescriptor(service.getApiServiceDescriptor());
StreamObserver<BeamFnApi.Elements> outboundObserver = BeamFnDataGrpc.newStub(channel).data(TestStreams.withOnNext(clientInboundElements::add).build());
waitForInboundElements.await();
outboundObserver.onCompleted();
return null;
});
}
for (int i = 0; i < 3; ++i) {
CloseableFnDataReceiver<WindowedValue<String>> consumer = service.getDataService(WORKER_ID).send(LogicalEndpoint.data(Integer.toString(i), TRANSFORM_ID), CODER);
consumer.accept(valueInGlobalWindow("A" + i));
consumer.accept(valueInGlobalWindow("B" + i));
consumer.accept(valueInGlobalWindow("C" + i));
consumer.close();
}
// Specifically copy the elements to a new list so we perform blocking calls on the queue
// to ensure the elements arrive.
List<Elements> copy = new ArrayList<>();
for (int i = 0; i < numberOfClients; ++i) {
copy.add(clientInboundElements.take());
}
assertThat(copy, containsInAnyOrder(elementsWithData("0"), elementsWithData("1"), elementsWithData("2")));
waitForInboundElements.countDown();
}
use of org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements.Data in project beam by apache.
the class BeamFnDataGrpcServiceTest method testMultipleClientsSendMessagesAreDirectedToProperConsumers.
@Test
public void testMultipleClientsSendMessagesAreDirectedToProperConsumers() throws Exception {
LinkedBlockingQueue<BeamFnApi.Elements> clientInboundElements = new LinkedBlockingQueue<>();
ExecutorService executorService = Executors.newCachedThreadPool();
CountDownLatch waitForInboundElements = new CountDownLatch(1);
for (int i = 0; i < 3; ++i) {
String instructionId = Integer.toString(i);
executorService.submit(() -> {
ManagedChannel channel = ManagedChannelFactory.createDefault().withInterceptors(Arrays.asList(AddHarnessIdInterceptor.create(WORKER_ID))).forDescriptor(service.getApiServiceDescriptor());
StreamObserver<BeamFnApi.Elements> outboundObserver = BeamFnDataGrpc.newStub(channel).data(TestStreams.withOnNext(clientInboundElements::add).build());
outboundObserver.onNext(elementsWithData(instructionId));
waitForInboundElements.await();
outboundObserver.onCompleted();
return null;
});
}
List<Collection<WindowedValue<String>>> serverInboundValues = new ArrayList<>();
Collection<InboundDataClient> inboundDataClients = new ArrayList<>();
for (int i = 0; i < 3; ++i) {
BlockingQueue<WindowedValue<String>> serverInboundValue = new LinkedBlockingQueue<>();
serverInboundValues.add(serverInboundValue);
inboundDataClients.add(service.getDataService(WORKER_ID).receive(LogicalEndpoint.data(Integer.toString(i), TRANSFORM_ID), CODER, serverInboundValue::add));
}
// Waiting for the client provides the necessary synchronization for the elements to arrive.
for (InboundDataClient inboundDataClient : inboundDataClients) {
inboundDataClient.awaitCompletion();
}
waitForInboundElements.countDown();
for (int i = 0; i < 3; ++i) {
assertThat(serverInboundValues.get(i), contains(valueInGlobalWindow("A" + i), valueInGlobalWindow("B" + i), valueInGlobalWindow("C" + i)));
}
assertThat(clientInboundElements, empty());
}
use of org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements.Data in project beam by apache.
the class BeamFnDataGrpcService method data.
@Override
public StreamObserver<Elements> data(final StreamObserver<Elements> outboundObserver) {
String sdkWorkerId = headerAccessor.getSdkWorkerId();
LOG.info("Beam Fn Data client connected for clientId {}", sdkWorkerId);
BeamFnDataGrpcMultiplexer multiplexer = new BeamFnDataGrpcMultiplexer(apiServiceDescriptor, OutboundObserverFactory.trivial(), (StreamObserver<BeamFnApi.Elements> inboundObserver) -> streamObserverFactory.apply(outboundObserver));
// First client that connects completes this future
getClientFuture(sdkWorkerId).complete(multiplexer);
try {
// incoming messages are sent to the single multiplexer instance.
return getClientFuture(sdkWorkerId).get().getInboundObserver();
} catch (InterruptedException | ExecutionException e) {
throw new RuntimeException(e);
}
}
use of org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements.Data in project beam by apache.
the class ProcessBundleHandlerTest method testDataProcessingExceptionsArePropagated.
@Test
public void testDataProcessingExceptionsArePropagated() throws Exception {
BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = BeamFnApi.ProcessBundleDescriptor.newBuilder().putTransforms("2L", RunnerApi.PTransform.newBuilder().setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()).build()).build();
Map<String, BeamFnApi.ProcessBundleDescriptor> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor);
Mockito.doAnswer((invocation) -> {
ByteString.Output encodedData = ByteString.newOutput();
StringUtf8Coder.of().encode("A", encodedData);
String instructionId = invocation.getArgument(0, String.class);
CloseableFnDataReceiver<BeamFnApi.Elements> data = invocation.getArgument(2, CloseableFnDataReceiver.class);
data.accept(BeamFnApi.Elements.newBuilder().addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(instructionId).setTransformId("2L").setData(encodedData.toByteString()).setIsLast(true)).build());
return null;
}).when(beamFnDataClient).registerReceiver(any(), any(), any());
ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), fnApiRegistry::get, beamFnDataClient, null, /* beamFnStateGrpcClientCache */
null, /* finalizeBundleHandler */
new ShortIdMap(), ImmutableMap.of(DATA_INPUT_URN, (PTransformRunnerFactory<Object>) (context) -> {
context.addIncomingDataEndpoint(ApiServiceDescriptor.getDefaultInstance(), StringUtf8Coder.of(), (input) -> {
throw new IllegalStateException("TestException");
});
return null;
}), Caches.noop(), new BundleProcessorCache());
assertThrows("TestException", IllegalStateException.class, () -> handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setInstructionId("instructionId").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build()));
// Ensure that we unregister during successful processing
verify(beamFnDataClient).registerReceiver(eq("instructionId"), any(), any());
verifyNoMoreInteractions(beamFnDataClient);
}
use of org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements.Data in project beam by apache.
the class ProcessBundleHandlerTest method testInstructionEmbeddedElementsAreProcessed.
@Test
public void testInstructionEmbeddedElementsAreProcessed() throws Exception {
List<String> dataOutput = new ArrayList<>();
List<Timers> timerOutput = new ArrayList<>();
ProcessBundleHandler handler = setupProcessBundleHandlerForSimpleRecordingDoFn(dataOutput, timerOutput, false);
ByteString.Output encodedData = ByteString.newOutput();
KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()).encode(KV.of("", "data"), encodedData);
ByteString.Output encodedTimer = ByteString.newOutput();
Timer.Coder.of(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE).encode(Timer.of("", "timer_id", Collections.singletonList(GlobalWindow.INSTANCE), Instant.ofEpochMilli(1L), Instant.ofEpochMilli(1L), PaneInfo.ON_TIME_AND_ONLY_FIRING), encodedTimer);
Elements elements = Elements.newBuilder().addData(Data.newBuilder().setInstructionId("998L").setTransformId("2L").setData(encodedData.toByteString()).build()).addData(Data.newBuilder().setInstructionId("998L").setTransformId("2L").setIsLast(true).build()).addTimers(Timers.newBuilder().setInstructionId("998L").setTransformId("3L").setTimerFamilyId(TimerFamilyDeclaration.PREFIX + SimpleDoFn.TIMER_FAMILY_ID).setTimers(encodedTimer.toByteString()).build()).addTimers(Timers.newBuilder().setInstructionId("998L").setTransformId("3L").setTimerFamilyId(TimerFamilyDeclaration.PREFIX + SimpleDoFn.TIMER_FAMILY_ID).setIsLast(true).build()).build();
handler.processBundle(InstructionRequest.newBuilder().setInstructionId("998L").setProcessBundle(ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L").setElements(elements)).build());
handler.shutdown();
assertThat(dataOutput, contains("data"));
Timer<String> timer = Timer.Coder.of(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE).decode(timerOutput.get(0).getTimers().newInput());
assertEquals("output_timer", timer.getDynamicTimerTag());
}
Aggregations