use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor in project beam by apache.
the class BeamFnMapTaskExecutorTest method testExtractCounterUpdatesReturnsValidProgressTrackerCounterUpdatesIfPresent.
@Test(timeout = ReadOperation.DEFAULT_PROGRESS_UPDATE_PERIOD_MS * 60)
public void testExtractCounterUpdatesReturnsValidProgressTrackerCounterUpdatesIfPresent() throws Exception {
final CountDownLatch progressSentLatch = new CountDownLatch(1);
final CountDownLatch processBundleLatch = new CountDownLatch(1);
final int expectedCounterValue = 5;
final MonitoringInfo expectedMonitoringInfo = MonitoringInfo.newBuilder().setUrn(MonitoringInfoConstants.Urns.USER_SUM_INT64).putLabels(MonitoringInfoConstants.Labels.NAME, "ExpectedCounter").putLabels(MonitoringInfoConstants.Labels.NAMESPACE, "anyString").setType(MonitoringInfoConstants.TypeUrns.SUM_INT64_TYPE).putLabels(MonitoringInfoConstants.Labels.PTRANSFORM, "ExpectedPTransform").setPayload(encodeInt64Counter(expectedCounterValue)).build();
InstructionRequestHandler instructionRequestHandler = new InstructionRequestHandler() {
@Override
public CompletionStage<InstructionResponse> handle(InstructionRequest request) {
switch(request.getRequestCase()) {
case REGISTER:
return CompletableFuture.completedFuture(responseFor(request).build());
case PROCESS_BUNDLE:
return MoreFutures.supplyAsync(() -> {
processBundleLatch.await();
return responseFor(request).setProcessBundle(BeamFnApi.ProcessBundleResponse.newBuilder().addMonitoringInfos(expectedMonitoringInfo)).build();
});
case PROCESS_BUNDLE_PROGRESS:
progressSentLatch.countDown();
return CompletableFuture.completedFuture(responseFor(request).setProcessBundleProgress(BeamFnApi.ProcessBundleProgressResponse.newBuilder().addMonitoringInfos(expectedMonitoringInfo)).build());
default:
throw new RuntimeException("Reached unexpected code path");
}
}
@Override
public void registerProcessBundleDescriptor(ProcessBundleDescriptor descriptor) {
}
@Override
public void close() {
}
};
Map<String, DataflowStepContext> stepContextMap = new HashMap<>();
stepContextMap.put("ExpectedPTransform", generateDataflowStepContext("Expected"));
RegisterAndProcessBundleOperation processOperation = new RegisterAndProcessBundleOperation(IdGenerators.decrementingLongs(), instructionRequestHandler, mockBeamFnStateDelegator, REGISTER_REQUEST, ImmutableMap.of(), stepContextMap, ImmutableMap.of(), ImmutableTable.of(), ImmutableMap.of(), mockContext);
BeamFnMapTaskExecutor mapTaskExecutor = BeamFnMapTaskExecutor.forOperations(ImmutableList.of(readOperation, grpcPortWriteOperation, processOperation), executionStateTracker);
// Launch the BeamFnMapTaskExecutor and wait until we are sure there has been one
// tentative update
CompletionStage<Void> doneFuture = MoreFutures.runAsync(mapTaskExecutor::execute);
progressSentLatch.await();
Iterable<CounterUpdate> metricsCounterUpdates = Collections.emptyList();
while (Iterables.size(metricsCounterUpdates) == 0) {
Thread.sleep(ReadOperation.DEFAULT_PROGRESS_UPDATE_PERIOD_MS);
metricsCounterUpdates = mapTaskExecutor.extractMetricUpdates();
}
// Get the final metrics
processBundleLatch.countDown();
MoreFutures.get(doneFuture);
metricsCounterUpdates = mapTaskExecutor.extractMetricUpdates();
assertThat(Iterables.size(metricsCounterUpdates), equalTo(1));
CounterUpdate resultCounter = metricsCounterUpdates.iterator().next();
assertTrue(new CounterHamcrestMatchers.CounterUpdateIntegerValueMatcher(expectedCounterValue).matches(resultCounter));
assertEquals("ExpectedCounter", resultCounter.getStructuredNameAndMetadata().getName().getName());
}
use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor in project beam by apache.
the class BeamFnControlService method getProcessBundleDescriptor.
@Override
public void getProcessBundleDescriptor(BeamFnApi.GetProcessBundleDescriptorRequest request, StreamObserver<BeamFnApi.ProcessBundleDescriptor> responseObserver) {
String bundleDescriptorId = request.getProcessBundleDescriptorId();
LOG.info("getProcessBundleDescriptor request with id {}", bundleDescriptorId);
BeamFnApi.ProcessBundleDescriptor descriptor = processBundleDescriptors.get(bundleDescriptorId);
if (descriptor == null) {
String msg = String.format("ProcessBundleDescriptor with id %s not found", bundleDescriptorId);
responseObserver.onError(new StatusException(Status.NOT_FOUND.withDescription(msg)));
LOG.error(msg);
} else {
responseObserver.onNext(descriptor);
responseObserver.onCompleted();
}
}
use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor in project beam by apache.
the class ProcessBundleHandlerTest method testPendingStateCallsBlockTillCompletion.
@Test
public void testPendingStateCallsBlockTillCompletion() throws Exception {
BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = BeamFnApi.ProcessBundleDescriptor.newBuilder().putTransforms("2L", RunnerApi.PTransform.newBuilder().setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()).build()).setStateApiServiceDescriptor(ApiServiceDescriptor.getDefaultInstance()).build();
Map<String, BeamFnApi.ProcessBundleDescriptor> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor);
CompletableFuture<StateResponse>[] successfulResponse = new CompletableFuture[1];
CompletableFuture<StateResponse>[] unsuccessfulResponse = new CompletableFuture[1];
BeamFnStateGrpcClientCache mockBeamFnStateGrpcClient = Mockito.mock(BeamFnStateGrpcClientCache.class);
BeamFnStateClient mockBeamFnStateClient = Mockito.mock(BeamFnStateClient.class);
when(mockBeamFnStateGrpcClient.forApiServiceDescriptor(any())).thenReturn(mockBeamFnStateClient);
doAnswer(invocation -> {
StateRequest.Builder stateRequestBuilder = (StateRequest.Builder) invocation.getArguments()[0];
CompletableFuture<StateResponse> completableFuture = new CompletableFuture<>();
new Thread(() -> {
// Simulate sleeping which introduces a race which most of the time requires
// the ProcessBundleHandler to block.
Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
switch(stateRequestBuilder.getInstructionId()) {
case "SUCCESS":
completableFuture.complete(StateResponse.getDefaultInstance());
break;
case "FAIL":
completableFuture.completeExceptionally(new RuntimeException("TEST ERROR"));
}
}).start();
return completableFuture;
}).when(mockBeamFnStateClient).handle(any());
ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), fnApiRegistry::get, beamFnDataClient, mockBeamFnStateGrpcClient, null, /* finalizeBundleHandler */
new ShortIdMap(), ImmutableMap.of(DATA_INPUT_URN, new PTransformRunnerFactory<Object>() {
@Override
public Object createRunnerForPTransform(Context context) throws IOException {
BeamFnStateClient beamFnStateClient = context.getBeamFnStateClient();
context.addStartBundleFunction(() -> doStateCalls(beamFnStateClient));
return null;
}
private void doStateCalls(BeamFnStateClient beamFnStateClient) {
successfulResponse[0] = beamFnStateClient.handle(StateRequest.newBuilder().setInstructionId("SUCCESS"));
unsuccessfulResponse[0] = beamFnStateClient.handle(StateRequest.newBuilder().setInstructionId("FAIL"));
}
}), Caches.noop(), new BundleProcessorCache());
handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build());
assertTrue(successfulResponse[0].isDone());
assertTrue(unsuccessfulResponse[0].isDone());
}
use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor in project beam by apache.
the class ProcessBundleHandlerTest method testTimerRegistrationsFailIfNoTimerApiServiceDescriptorSpecified.
@Test
public void testTimerRegistrationsFailIfNoTimerApiServiceDescriptorSpecified() 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);
ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), fnApiRegistry::get, beamFnDataClient, null, /* beamFnStateGrpcClientCache */
null, /* finalizeBundleHandler */
new ShortIdMap(), ImmutableMap.of(DATA_INPUT_URN, new PTransformRunnerFactory<Object>() {
@Override
public Object createRunnerForPTransform(Context context) throws IOException {
context.addOutgoingTimersEndpoint("timer", Timer.Coder.of(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE));
return null;
}
}), Caches.noop(), new BundleProcessorCache());
assertThrows("Timers are unsupported", IllegalStateException.class, () -> handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build()));
}
use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor 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);
}
Aggregations