public BeamFnApi.InstructionResponse.Builder processBundle(BeamFnApi.InstructionRequest request)
      throws Exception {
    BeamFnApi.ProcessBundleResponse.Builder response = BeamFnApi.ProcessBundleResponse.newBuilder();

    BundleProcessor bundleProcessor =
        bundleProcessorCache.get(
            request,
            () -> {
              try {
                return createBundleProcessor(
                    request.getProcessBundle().getProcessBundleDescriptorId(),
                    request.getProcessBundle());
              } catch (IOException e) {
                throw new RuntimeException(e);
              }
            });
    try {
      PTransformFunctionRegistry startFunctionRegistry = bundleProcessor.getStartFunctionRegistry();
      PTransformFunctionRegistry finishFunctionRegistry =
          bundleProcessor.getFinishFunctionRegistry();
      ExecutionStateTracker stateTracker = bundleProcessor.getStateTracker();

      try (HandleStateCallsForBundle beamFnStateClient = bundleProcessor.getBeamFnStateClient()) {
        stateTracker.start(request.getInstructionId());
        try {
          // Already in reverse topological order so we don't need to do anything.
          for (ThrowingRunnable startFunction : startFunctionRegistry.getFunctions()) {
            LOG.debug("Starting function {}", startFunction);
            startFunction.run();
          }

          if (request.getProcessBundle().hasElements()) {
            boolean inputFinished =
                bundleProcessor
                    .getInboundObserver()
                    .multiplexElements(request.getProcessBundle().getElements());
            if (!inputFinished) {
              throw new RuntimeException(
                  "Elements embedded in ProcessBundleRequest do not contain stream terminators for "
                      + "all data and timer inputs. Unterminated endpoints: "
                      + bundleProcessor.getInboundObserver().getUnfinishedEndpoints());
            }
          } else if (!bundleProcessor.getInboundEndpointApiServiceDescriptors().isEmpty()) {
            BeamFnDataInboundObserver observer = bundleProcessor.getInboundObserver();
            beamFnDataClient.registerReceiver(
                request.getInstructionId(),
                bundleProcessor.getInboundEndpointApiServiceDescriptors(),
                observer);
            observer.awaitCompletion();
            beamFnDataClient.unregisterReceiver(
                request.getInstructionId(),
                bundleProcessor.getInboundEndpointApiServiceDescriptors());
          }

          // Need to reverse this since we want to call finish in topological order.
          for (ThrowingRunnable finishFunction :
              Lists.reverse(finishFunctionRegistry.getFunctions())) {
            LOG.debug("Finishing function {}", finishFunction);
            finishFunction.run();
          }

          // If bundleProcessor has not flushed any elements, embed them in response.
          embedOutboundElementsIfApplicable(response, bundleProcessor);

          // Add all checkpointed residuals to the response.
          response.addAllResidualRoots(bundleProcessor.getSplitListener().getResidualRoots());

          // Add all metrics to the response.
          bundleProcessor.getProgressRequestLock().lock();
          Map<String, ByteString> monitoringData = finalMonitoringData(bundleProcessor);
          if (runnerAcceptsShortIds) {
            response.putAllMonitoringData(monitoringData);
          } else {
            for (Map.Entry<String, ByteString> metric : monitoringData.entrySet()) {
              response.addMonitoringInfos(
                  shortIds.get(metric.getKey()).toBuilder().setPayload(metric.getValue()));
            }
          }

          if (!bundleProcessor.getBundleFinalizationCallbackRegistrations().isEmpty()) {
            finalizeBundleHandler.registerCallbacks(
                bundleProcessor.getInstructionId(),
                ImmutableList.copyOf(bundleProcessor.getBundleFinalizationCallbackRegistrations()));
            response.setRequiresFinalization(true);
          }
        } finally {
          // We specifically deactivate state tracking while we are holding the progress request and
          // sampling locks.
          stateTracker.reset();
        }
      }

      // Mark the bundle processor as re-usable.
      bundleProcessorCache.release(
          request.getProcessBundle().getProcessBundleDescriptorId(), bundleProcessor);
      return BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response);
    } catch (Exception e) {
      // Make sure we clean-up from the active set of bundle processors.
      bundleProcessorCache.discard(bundleProcessor);
      throw e;
    }
  }
@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(),
            executionStateSampler,
            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(),
            null /* dataSampler */);
    handler.processBundle(
        BeamFnApi.InstructionRequest.newBuilder()
            .setProcessBundle(
                BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L"))
            .build());

    assertTrue(successfulResponse[0].isDone());
    assertTrue(unsuccessfulResponse[0].isDone());
  }