RegisterAndProcessBundleOperation extends Operation { @Override public void start() throws Exception { try (Closeable scope = context.enterStart()) { super.start(); if (registerFuture == null) { InstructionRequest request = InstructionRequest.newBuilder() .setInstructionId(idGenerator.getId()) .setRegister(registerRequest) .build(); registerFuture = instructionRequestHandler.handle(request); } checkState( registerRequest.getProcessBundleDescriptorCount() == 1, "Only one bundle registration at a time currently supported."); InstructionRequest processBundleRequest = InstructionRequest.newBuilder() .setInstructionId(getProcessBundleInstructionId()) .setProcessBundle( ProcessBundleRequest.newBuilder() .setProcessBundleDescriptorId( registerRequest.getProcessBundleDescriptor(0).getId())) .build(); deregisterStateHandler = beamFnStateDelegator.registerForProcessBundleInstructionId( getProcessBundleInstructionId(), this::delegateByStateKeyType); processBundleResponse = getRegisterResponse(registerFuture) .thenCompose( registerResponse -> instructionRequestHandler.handle(processBundleRequest)); } } RegisterAndProcessBundleOperation(
      IdGenerator idGenerator,
      InstructionRequestHandler instructionRequestHandler,
      StateDelegator beamFnStateDelegator,
      RegisterRequest registerRequest,
      Map<String, DataflowOperationContext> ptransformIdToOperationContext,
      Map<String, DataflowStepContext> ptransformIdToSystemStepContext,
      Map<String, SideInputReader> ptransformIdToSideInputReader,
      Table<String, String, PCollectionView<?>> ptransformIdToSideInputIdToPCollectionView,
      Map<String, NameContext> pcollectionIdToNameContext,
      OperationContext context); synchronized String getProcessBundleInstructionId(); String getCurrentProcessBundleInstructionId(); @Override void start(); @Override void finish(); @Override void abort(); Map<String, NameContext> getPCollectionIdToNameContext(); Map<String, DataflowStepContext> getPtransformIdToUserStepContext(); CompletionStage<BeamFnApi.ProcessBundleProgressResponse> getProcessBundleProgress(); CompletionStage<List<MonitoringInfo>> getFinalMonitoringInfos(); boolean hasFailed(); List<MonitoringInfo> findIOPCollectionMonitoringInfos(
      Iterable<MonitoringInfo> monitoringInfos); @Override boolean supportsRestart(); @Override String toString();  }
@Test public void testAbortCancelsAndCleansUpDuringRegister() throws Exception { IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L); ExecutorService executorService = Executors.newCachedThreadPool(); CountDownLatch waitForAbortToComplete = new CountDownLatch(1); AtomicReference<ThrowingRunnable> abortReference = new AtomicReference<>(); RegisterAndProcessBundleOperation operation = new RegisterAndProcessBundleOperation( idGenerator, new TestInstructionRequestHandler() { @Override public CompletionStage<InstructionResponse> handle(InstructionRequest request) { CompletableFuture<InstructionResponse> responseFuture = new CompletableFuture<>(); if (request.getRequestCase() == RequestCase.PROCESS_BUNDLE) { executorService.submit( (Callable<Void>) () -> { abortReference.get().run(); waitForAbortToComplete.countDown(); return null; }); } else { completeFuture(request, responseFuture); } return responseFuture; } }, mockBeamFnStateDelegator, REGISTER_REQUEST, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(), ImmutableTable.of(), ImmutableMap.of(), mockContext); abortReference.set(operation::abort); operation.start(); waitForAbortToComplete.await(); assertEquals(stateServiceRegisterCounter.get(), stateServiceAbortCounter.get()); assertEquals(0, stateServiceDeregisterCounter.get()); }
@Test public void testAbortCancelsAndCleansUpDuringProcessBundle() throws Exception { IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L); ExecutorService executorService = Executors.newCachedThreadPool(); CountDownLatch waitForAbortToComplete = new CountDownLatch(1); AtomicReference<ThrowingRunnable> abortReference = new AtomicReference<>(); RegisterAndProcessBundleOperation operation = new RegisterAndProcessBundleOperation( idGenerator, new TestInstructionRequestHandler() { @Override public CompletionStage<InstructionResponse> handle(InstructionRequest request) { CompletableFuture<InstructionResponse> responseFuture = new CompletableFuture<>(); if (request.getRequestCase() == RequestCase.PROCESS_BUNDLE) { executorService.submit( (Callable<Void>) () -> { abortReference.get().run(); waitForAbortToComplete.countDown(); return null; }); } else { completeFuture(request, responseFuture); } return responseFuture; } }, mockBeamFnStateDelegator, REGISTER_REQUEST, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(), ImmutableTable.of(), ImmutableMap.of(), mockContext); abortReference.set(operation::abort); operation.start(); waitForAbortToComplete.await(); assertEquals(stateServiceRegisterCounter.get(), stateServiceAbortCounter.get()); assertEquals(0, stateServiceDeregisterCounter.get()); }