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 testInstructionIsUnregisteredFromBeamFnDataClientOnSuccess() 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) -> {
              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")
                              .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(),
            executionStateSampler,
            ImmutableMap.of(
                DATA_INPUT_URN,
                (PTransformRunnerFactory<Object>)
                    (context) -> {
                      context.addIncomingDataEndpoint(
                          ApiServiceDescriptor.getDefaultInstance(),
                          StringUtf8Coder.of(),
                          (input) -> {});
                      return null;
                    }),
            Caches.noop(),
            new BundleProcessorCache(),
            null /* dataSampler */);
    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());
    verify(beamFnDataClient).unregisterReceiver(eq("instructionId"), any());
    verifyNoMoreInteractions(beamFnDataClient);
  }