ProcessBundleHandler { public BeamFnApi.InstructionResponse.Builder processBundle(BeamFnApi.InstructionRequest request) throws Exception { BeamFnApi.ProcessBundleResponse.Builder response = BeamFnApi.ProcessBundleResponse.newBuilder(); BundleProcessor bundleProcessor = bundleProcessorCache.get( request.getProcessBundle().getProcessBundleDescriptorId(), request.getInstructionId(), () -> { try { return createBundleProcessor( request.getProcessBundle().getProcessBundleDescriptorId(), request.getProcessBundle()); } catch (IOException e) { throw new RuntimeException(e); } }); PTransformFunctionRegistry startFunctionRegistry = bundleProcessor.getStartFunctionRegistry(); PTransformFunctionRegistry finishFunctionRegistry = bundleProcessor.getFinishFunctionRegistry(); ExecutionStateTracker stateTracker = bundleProcessor.getStateTracker(); QueueingBeamFnDataClient queueingClient = bundleProcessor.getQueueingClient(); try (HandleStateCallsForBundle beamFnStateClient = bundleProcessor.getBeamFnStateClient()) { try (Closeable closeTracker = stateTracker.activate()) { for (ThrowingRunnable startFunction : startFunctionRegistry.getFunctions()) { LOG.debug("Starting function {}", startFunction); startFunction.run(); } queueingClient.drainAndBlock(); for (ThrowingRunnable finishFunction : Lists.reverse(finishFunctionRegistry.getFunctions())) { LOG.debug("Finishing function {}", finishFunction); finishFunction.run(); } } response.addAllResidualRoots(bundleProcessor.getSplitListener().getResidualRoots()); response.addAllMonitoringInfos( bundleProcessor.getStartFunctionRegistry().getExecutionTimeMonitoringInfos()); response.addAllMonitoringInfos( bundleProcessor.getpCollectionConsumerRegistry().getExecutionTimeMonitoringInfos()); response.addAllMonitoringInfos( bundleProcessor.getFinishFunctionRegistry().getExecutionTimeMonitoringInfos()); response.addAllMonitoringInfos( bundleProcessor.getMetricsContainerRegistry().getMonitoringInfos()); for (ProgressRequestCallback progressRequestCallback : bundleProcessor.getProgressRequestCallbacks()) { response.addAllMonitoringInfos(progressRequestCallback.getMonitoringInfos()); } if (!bundleProcessor.getBundleFinalizationCallbackRegistrations().isEmpty()) { finalizeBundleHandler.registerCallbacks( bundleProcessor.getInstructionId(), ImmutableList.copyOf(bundleProcessor.getBundleFinalizationCallbackRegistrations())); response.setRequiresFinalization(true); } bundleProcessorCache.release( request.getProcessBundle().getProcessBundleDescriptorId(), bundleProcessor); } return BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response); } ProcessBundleHandler(
      PipelineOptions options,
      Function<String, Message> fnApiRegistry,
      BeamFnDataClient beamFnDataClient,
      BeamFnStateGrpcClientCache beamFnStateGrpcClientCache,
      FinalizeBundleHandler finalizeBundleHandler); @VisibleForTesting  ProcessBundleHandler(
      PipelineOptions options,
      Function<String, Message> fnApiRegistry,
      BeamFnDataClient beamFnDataClient,
      BeamFnStateGrpcClientCache beamFnStateGrpcClientCache,
      FinalizeBundleHandler finalizeBundleHandler,
      Map<String, PTransformRunnerFactory> urnToPTransformRunnerFactoryMap,
      BundleProcessorCache bundleProcessorCache); BeamFnApi.InstructionResponse.Builder processBundle(BeamFnApi.InstructionRequest request); BeamFnApi.InstructionResponse.Builder progress(BeamFnApi.InstructionRequest request); BeamFnApi.InstructionResponse.Builder trySplit(BeamFnApi.InstructionRequest request); void shutdown(); static final String JAVA_SOURCE_URN; }
@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, Message> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor); List<RunnerApi.PTransform> transformsProcessed = new ArrayList<>(); List<String> orderOfOperations = new ArrayList<>(); PTransformRunnerFactory<Object> startFinishRecorder = (pipelineOptions, beamFnDataClient, beamFnStateClient, beamFnTimerClient, pTransformId, pTransform, processBundleInstructionId, pCollections, coders, windowingStrategies, pCollectionConsumerRegistry, startFunctionRegistry, finishFunctionRegistry, addResetFunction, addTearDownFunction, addProgressRequestCallback, splitListener, bundleFinalizer) -> { transformsProcessed.add(pTransform); startFunctionRegistry.register( pTransformId, () -> { assertThat(processBundleInstructionId.get(), equalTo("999L")); orderOfOperations.add("Start" + pTransformId); }); finishFunctionRegistry.register( pTransformId, () -> { assertThat(processBundleInstructionId.get(), equalTo("999L")); orderOfOperations.add("Finish" + pTransformId); }); return null; }; ProcessBundleHandler handler = new ProcessBundleHandler( PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient, null , null , ImmutableMap.of( DATA_INPUT_URN, startFinishRecorder, DATA_OUTPUT_URN, startFinishRecorder), new BundleProcessorCache()); handler.processBundle( BeamFnApi.InstructionRequest.newBuilder() .setInstructionId("999L") .setProcessBundle( BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")) .build()); assertThat( transformsProcessed, contains( processBundleDescriptor.getTransformsMap().get("3L"), processBundleDescriptor.getTransformsMap().get("2L"))); assertThat(orderOfOperations, contains("Start3L", "Start2L", "Finish2L", "Finish3L")); }
@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, Message> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor); ProcessBundleHandler handler = new ProcessBundleHandler( PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient, null , null , ImmutableMap.of( DATA_INPUT_URN, (pipelineOptions, beamFnDataClient, beamFnStateClient, beamFnTimerClient, pTransformId, pTransform, processBundleInstructionId, pCollections, coders, windowingStrategies, pCollectionConsumerRegistry, startFunctionRegistry, finishFunctionRegistry, addResetFunction, addTearDownFunction, addProgressRequestCallback, splitListener, bundleFinalizer) -> null), new TestBundleProcessorCache()); assertThat(TestBundleProcessor.resetCnt, equalTo(0)); handler.processBundle( BeamFnApi.InstructionRequest.newBuilder() .setInstructionId("998L") .setProcessBundle( BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")) .build()); assertThat(TestBundleProcessor.resetCnt, equalTo(1)); assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().size(), equalTo(1)); assertThat( handler.bundleProcessorCache.getCachedBundleProcessors().get("1L").size(), equalTo(1)); 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()); assertThat( handler.bundleProcessorCache.getCachedBundleProcessors().get("1L").size(), equalTo(0)); }
@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, Message> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor); ProcessBundleHandler handler = new ProcessBundleHandler( PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient, null , null , ImmutableMap.of( DATA_INPUT_URN, (pipelineOptions, beamFnDataClient, beamFnStateClient, beamFnTimerClient, pTransformId, pTransform, processBundleInstructionId, pCollections, coders, windowingStrategies, pCollectionConsumerRegistry, startFunctionRegistry, finishFunctionRegistry, addResetFunction, addTearDownFunction, addProgressRequestCallback, splitListener, bundleFinalizer) -> { thrown.expect(IllegalStateException.class); thrown.expectMessage("TestException"); throw new IllegalStateException("TestException"); }), new BundleProcessorCache()); handler.processBundle( BeamFnApi.InstructionRequest.newBuilder() .setProcessBundle( BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")) .build()); }
@Test public void testBundleFinalizationIsPropagated() 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, Message> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor); FinalizeBundleHandler mockFinalizeBundleHandler = mock(FinalizeBundleHandler.class); BundleFinalizer.Callback mockCallback = mock(BundleFinalizer.Callback.class); ProcessBundleHandler handler = new ProcessBundleHandler( PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient, null , mockFinalizeBundleHandler, ImmutableMap.of( DATA_INPUT_URN, (PTransformRunnerFactory<Object>) (pipelineOptions, beamFnDataClient, beamFnStateClient, beamFnTimerClient, pTransformId, pTransform, processBundleInstructionId, pCollections, coders, windowingStrategies, pCollectionConsumerRegistry, startFunctionRegistry, finishFunctionRegistry, addResetFunction, addTearDownFunction, addProgressRequestCallback, splitListener, bundleFinalizer) -> { startFunctionRegistry.register( pTransformId, () -> bundleFinalizer.afterBundleCommit( Instant.ofEpochMilli(42L), mockCallback)); return null; }), new BundleProcessorCache()); BeamFnApi.InstructionResponse.Builder response = handler.processBundle( BeamFnApi.InstructionRequest.newBuilder() .setInstructionId("2L") .setProcessBundle( BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")) .build()); assertTrue(response.getProcessBundle().getRequiresFinalization()); verify(mockFinalizeBundleHandler) .registerCallbacks( eq("2L"), argThat( (Collection<CallbackRegistration> arg) -> { CallbackRegistration registration = Iterables.getOnlyElement(arg); assertEquals(Instant.ofEpochMilli(42L), registration.getExpiryTime()); assertSame(mockCallback, registration.getCallback()); return true; })); }
@Test public void testPTransformStartExceptionsArePropagated() 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, Message> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor); ProcessBundleHandler handler = new ProcessBundleHandler( PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient, null , null , ImmutableMap.of( DATA_INPUT_URN, (PTransformRunnerFactory<Object>) (pipelineOptions, beamFnDataClient, beamFnStateClient, beamFnTimerClient, pTransformId, pTransform, processBundleInstructionId, pCollections, coders, windowingStrategies, pCollectionConsumerRegistry, startFunctionRegistry, finishFunctionRegistry, addResetFunction, addTearDownFunction, addProgressRequestCallback, splitListener, bundleFinalizer) -> { thrown.expect(IllegalStateException.class); thrown.expectMessage("TestException"); startFunctionRegistry.register( pTransformId, ProcessBundleHandlerTest::throwException); return null; }), new BundleProcessorCache()); handler.processBundle( BeamFnApi.InstructionRequest.newBuilder() .setProcessBundle( BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")) .build()); assertThat( handler.bundleProcessorCache.getCachedBundleProcessors(), equalTo(Collections.EMPTY_MAP)); }
@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, Message> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor); ProcessBundleHandler handler = new ProcessBundleHandler( PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient, null , null , ImmutableMap.of( DATA_INPUT_URN, (PTransformRunnerFactory<Object>) (pipelineOptions, beamFnDataClient, beamFnStateClient, beamFnTimerClient, pTransformId, pTransform, processBundleInstructionId, pCollections, coders, windowingStrategies, pCollectionConsumerRegistry, startFunctionRegistry, finishFunctionRegistry, addResetFunction, addTearDownFunction, addProgressRequestCallback, splitListener, bundleFinalizer) -> { thrown.expect(IllegalStateException.class); thrown.expectMessage("TestException"); finishFunctionRegistry.register( pTransformId, ProcessBundleHandlerTest::throwException); return null; }), new BundleProcessorCache()); handler.processBundle( BeamFnApi.InstructionRequest.newBuilder() .setProcessBundle( BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")) .build()); assertThat( handler.bundleProcessorCache.getCachedBundleProcessors(), equalTo(Collections.EMPTY_MAP)); }
@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, Message> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor); CompletableFuture<StateResponse> successfulResponse = new CompletableFuture<>(); CompletableFuture<StateResponse> unsuccessfulResponse = new CompletableFuture<>(); 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 = (CompletableFuture<StateResponse>) invocation.getArguments()[1]; new Thread( () -> { 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 null; }) .when(mockBeamFnStateClient) .handle(any(), any()); ProcessBundleHandler handler = new ProcessBundleHandler( PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient, mockBeamFnStateGrpcClient, null , ImmutableMap.of( DATA_INPUT_URN, new PTransformRunnerFactory<Object>() { @Override public Object createRunnerForPTransform( PipelineOptions pipelineOptions, BeamFnDataClient beamFnDataClient, BeamFnStateClient beamFnStateClient, BeamFnTimerClient beamFnTimerClient, String pTransformId, PTransform pTransform, Supplier<String> processBundleInstructionId, Map<String, PCollection> pCollections, Map<String, Coder> coders, Map<String, WindowingStrategy> windowingStrategies, PCollectionConsumerRegistry pCollectionConsumerRegistry, PTransformFunctionRegistry startFunctionRegistry, PTransformFunctionRegistry finishFunctionRegistry, Consumer<ThrowingRunnable> addResetFunction, Consumer<ThrowingRunnable> addTearDownFunction, Consumer<ProgressRequestCallback> addProgressRequestCallback, BundleSplitListener splitListener, BundleFinalizer bundleFinalizer) throws IOException { startFunctionRegistry.register( pTransformId, () -> doStateCalls(beamFnStateClient)); return null; } private void doStateCalls(BeamFnStateClient beamFnStateClient) { beamFnStateClient.handle( StateRequest.newBuilder().setInstructionId("SUCCESS"), successfulResponse); beamFnStateClient.handle( StateRequest.newBuilder().setInstructionId("FAIL"), unsuccessfulResponse); } }), new BundleProcessorCache()); handler.processBundle( BeamFnApi.InstructionRequest.newBuilder() .setProcessBundle( BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")) .build()); assertTrue(successfulResponse.isDone()); assertTrue(unsuccessfulResponse.isDone()); }
@Test public void testStateCallsFailIfNoStateApiServiceDescriptorSpecified() 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, Message> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor); ProcessBundleHandler handler = new ProcessBundleHandler( PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient, null , null , ImmutableMap.of( DATA_INPUT_URN, new PTransformRunnerFactory<Object>() { @Override public Object createRunnerForPTransform( PipelineOptions pipelineOptions, BeamFnDataClient beamFnDataClient, BeamFnStateClient beamFnStateClient, BeamFnTimerClient beamFnTimerClient, String pTransformId, PTransform pTransform, Supplier<String> processBundleInstructionId, Map<String, PCollection> pCollections, Map<String, Coder> coders, Map<String, WindowingStrategy> windowingStrategies, PCollectionConsumerRegistry pCollectionConsumerRegistry, PTransformFunctionRegistry startFunctionRegistry, PTransformFunctionRegistry finishFunctionRegistry, Consumer<ThrowingRunnable> addResetFunction, Consumer<ThrowingRunnable> addTearDownFunction, Consumer<ProgressRequestCallback> addProgressRequestCallback, BundleSplitListener splitListener, BundleFinalizer bundleFinalizer) throws IOException { startFunctionRegistry.register( pTransformId, () -> doStateCalls(beamFnStateClient)); return null; } private void doStateCalls(BeamFnStateClient beamFnStateClient) { thrown.expect(IllegalStateException.class); thrown.expectMessage("State API calls are unsupported"); beamFnStateClient.handle( StateRequest.newBuilder().setInstructionId("SUCCESS"), new CompletableFuture<>()); } }), new BundleProcessorCache()); handler.processBundle( BeamFnApi.InstructionRequest.newBuilder() .setProcessBundle( BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")) .build()); }
@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, Message> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor); ProcessBundleHandler handler = new ProcessBundleHandler( PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient, null , null , ImmutableMap.of( DATA_INPUT_URN, new PTransformRunnerFactory<Object>() { @Override public Object createRunnerForPTransform( PipelineOptions pipelineOptions, BeamFnDataClient beamFnDataClient, BeamFnStateClient beamFnStateClient, BeamFnTimerClient beamFnTimerClient, String pTransformId, PTransform pTransform, Supplier<String> processBundleInstructionId, Map<String, PCollection> pCollections, Map<String, Coder> coders, Map<String, WindowingStrategy> windowingStrategies, PCollectionConsumerRegistry pCollectionConsumerRegistry, PTransformFunctionRegistry startFunctionRegistry, PTransformFunctionRegistry finishFunctionRegistry, Consumer<ThrowingRunnable> addResetFunction, Consumer<ThrowingRunnable> addTearDownFunction, Consumer<ProgressRequestCallback> addProgressRequestCallback, BundleSplitListener splitListener, BundleFinalizer bundleFinalizer) throws IOException { startFunctionRegistry.register( pTransformId, () -> doTimerRegistrations(beamFnTimerClient)); return null; } private void doTimerRegistrations(BeamFnTimerClient beamFnTimerClient) { thrown.expect(IllegalStateException.class); thrown.expectMessage("Timers are unsupported"); beamFnTimerClient.register( LogicalEndpoint.timer("1L", "2L", "Timer"), Timer.Coder.of(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE), (timer) -> {}); } }), new BundleProcessorCache()); handler.processBundle( BeamFnApi.InstructionRequest.newBuilder() .setProcessBundle( BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")) .build()); }