use of org.apache.beam.runners.core.metrics.ShortIdMap in project beam by apache.
the class ProcessBundleHandlerTest method testCreatingPTransformExceptionsArePropagated.
@Test
public void testCreatingPTransformExceptionsArePropagated() 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, (context) -> {
throw new IllegalStateException("TestException");
}), Caches.noop(), new BundleProcessorCache());
assertThrows("TestException", IllegalStateException.class, () -> handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build()));
}
use of org.apache.beam.runners.core.metrics.ShortIdMap in project beam by apache.
the class ProcessBundleHandlerTest method testPTransformStartExceptionsArePropagated.
@Test
public void testPTransformStartExceptionsArePropagated() {
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, (PTransformRunnerFactory<Object>) (context) -> {
context.addStartBundleFunction(ProcessBundleHandlerTest::throwException);
return null;
}), Caches.noop(), new BundleProcessorCache());
assertThrows("TestException", IllegalStateException.class, () -> handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build()));
// BundleProcessor is not re-added back to the BundleProcessorCache in case of an exception
// during bundle processing
assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().get("1L"), empty());
}
use of org.apache.beam.runners.core.metrics.ShortIdMap in project beam by apache.
the class ProcessBundleHandlerTest method testPTransformFinishExceptionsArePropagated.
@Test
public void testPTransformFinishExceptionsArePropagated() 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, (PTransformRunnerFactory<Object>) (context) -> {
context.addFinishBundleFunction(ProcessBundleHandlerTest::throwException);
return null;
}), Caches.noop(), new BundleProcessorCache());
assertThrows("TestException", IllegalStateException.class, () -> handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build()));
// BundleProcessor is not re-added back to the BundleProcessorCache in case of an exception
// during bundle processing
assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().get("1L"), empty());
}
use of org.apache.beam.runners.core.metrics.ShortIdMap in project beam by apache.
the class ProcessBundleHandlerTest method testOrderOfStartAndFinishCalls.
@Test
public void testOrderOfStartAndFinishCalls() throws Exception {
BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = BeamFnApi.ProcessBundleDescriptor.newBuilder().putTransforms("2L", RunnerApi.PTransform.newBuilder().setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()).putOutputs("2L-output", "2L-output-pc").build()).putTransforms("3L", RunnerApi.PTransform.newBuilder().setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_OUTPUT_URN).build()).putInputs("3L-input", "2L-output-pc").build()).putPcollections("2L-output-pc", RunnerApi.PCollection.getDefaultInstance()).build();
Map<String, BeamFnApi.ProcessBundleDescriptor> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor);
List<RunnerApi.PTransform> transformsProcessed = new ArrayList<>();
List<String> orderOfOperations = new ArrayList<>();
PTransformRunnerFactory<Object> startFinishRecorder = (context) -> {
String pTransformId = context.getPTransformId();
transformsProcessed.add(context.getPTransform());
Supplier<String> processBundleInstructionId = context.getProcessBundleInstructionIdSupplier();
context.addStartBundleFunction(() -> {
assertThat(processBundleInstructionId.get(), equalTo("999L"));
orderOfOperations.add("Start" + pTransformId);
});
context.addFinishBundleFunction(() -> {
assertThat(processBundleInstructionId.get(), equalTo("999L"));
orderOfOperations.add("Finish" + pTransformId);
});
return null;
};
ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), fnApiRegistry::get, beamFnDataClient, null, /* beamFnStateClient */
null, /* finalizeBundleHandler */
new ShortIdMap(), ImmutableMap.of(DATA_INPUT_URN, startFinishRecorder, DATA_OUTPUT_URN, startFinishRecorder), Caches.noop(), new BundleProcessorCache());
handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setInstructionId("999L").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build());
// Processing of transforms is performed in reverse order.
assertThat(transformsProcessed, contains(processBundleDescriptor.getTransformsMap().get("3L"), processBundleDescriptor.getTransformsMap().get("2L")));
// Start should occur in reverse order while finish calls should occur in forward order
assertThat(orderOfOperations, contains("Start3L", "Start2L", "Finish2L", "Finish3L"));
}
use of org.apache.beam.runners.core.metrics.ShortIdMap in project beam by apache.
the class ProcessBundleHandlerTest method testBundleProcessorIsResetWhenAddedBackToCache.
@Test
public void testBundleProcessorIsResetWhenAddedBackToCache() 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, (context) -> null), Caches.noop(), new TestBundleProcessorCache());
assertThat(TestBundleProcessor.resetCnt, equalTo(0));
handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setInstructionId("998L").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build());
// Check that BundleProcessor is reset when added back to the cache
assertThat(TestBundleProcessor.resetCnt, equalTo(1));
// BundleProcessor is added back to the BundleProcessorCache
assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().size(), equalTo(1));
assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().get("1L").size(), equalTo(1));
// Add a reset handler that throws to test discarding the bundle processor on reset failure.
Iterables.getOnlyElement(handler.bundleProcessorCache.getCachedBundleProcessors().get("1L")).getResetFunctions().add(() -> {
throw new IllegalStateException("ResetFailed");
});
handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setInstructionId("999L").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build());
// BundleProcessor is discarded instead of being added back to the BundleProcessorCache
assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().get("1L").size(), equalTo(0));
}
Aggregations