@Override
  public StreamObserver<BeamFnApi.InstructionResponse> control(
      StreamObserver<BeamFnApi.InstructionRequest> requestObserver) {
    final String workerId = headerAccessor.getSdkWorkerId();
    LOG.info("Beam Fn Control client connected with id {}", workerId);
    FnApiControlClient newClient =
        FnApiControlClient.forRequestObserver(workerId, requestObserver, processBundleDescriptors);
    try {
      // Add the client to the pool of vended clients before making it available - we should close
      // the client when we close even if no one has picked it up yet. This can occur after the
      // service is closed, in which case the client will be discarded when the service is
      // discarded, which should be performed by a call to #shutdownNow. The remote caller must be
      // able to handle an unexpectedly terminated connection.
      synchronized (lock) {
        checkState(
            !closed, "%s already closed", FnApiControlClientPoolService.class.getSimpleName());
        // TODO: https://github.com/apache/beam/issues/18790: Prevent stale client references
        // from leaking.
        vendedClients.put(workerId, newClient);
      }
      // We do not attempt to transactionally add the client to our internal list and offer it to
      // the sink.
      clientSink.put(headerAccessor.getSdkWorkerId(), newClient);
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
    return newClient.asResponseObserver();
  }
@Test
  public void testIncomingConnection() throws Exception {
    String id = "fakeInstruction";
    CompletableFuture<StreamObserver<InstructionResponse>> clientResponseStream =
        new CompletableFuture<>();
    clientResponseStream.complete(
        stub.control(
            TestStreams.<InstructionRequest>withOnNext(
                    (request) -> {
                      try {
                        clientResponseStream
                            .get()
                            .onNext(
                                InstructionResponse.newBuilder()
                                    .setInstructionId(request.getInstructionId())
                                    .build());
                      } catch (Exception e) {
                        fail("Unexpected failure");
                        throw new RuntimeException(e);
                      }
                    })
                .build()));

    InstructionRequestHandler client = pool.getSource().take(WORKER_ID, Duration.ofSeconds(2));
    // Check that the client is wired up to the request channel
    CompletionStage<BeamFnApi.InstructionResponse> responseFuture =
        client.handle(BeamFnApi.InstructionRequest.newBuilder().setInstructionId(id).build());
    assertEquals(id, MoreFutures.get(responseFuture).getInstructionId());
  }