public void testThreadContextRestored() throws Exception {
   String header = randomAsciiOfLength(5);
   threadPool.getThreadContext().putHeader("test", header);
   AtomicBoolean called = new AtomicBoolean();
   sourceWithMockedRemoteCall("start_ok.json")
       .doStart(
           r -> {
             assertEquals(header, threadPool.getThreadContext().getHeader("test"));
             called.set(true);
           });
   assertTrue(called.get());
 }
 @Inject
 public LocalTransport(
     Settings settings,
     ThreadPool threadPool,
     Version version,
     NamedWriteableRegistry namedWriteableRegistry,
     CircuitBreakerService circuitBreakerService) {
   super(settings);
   this.threadPool = threadPool;
   this.version = version;
   int workerCount =
       this.settings.getAsInt(
           TRANSPORT_LOCAL_WORKERS, EsExecutors.boundedNumberOfProcessors(settings));
   int queueSize = this.settings.getAsInt(TRANSPORT_LOCAL_QUEUE, -1);
   logger.debug("creating [{}] workers, queue_size [{}]", workerCount, queueSize);
   final ThreadFactory threadFactory =
       EsExecutors.daemonThreadFactory(this.settings, LOCAL_TRANSPORT_THREAD_NAME_PREFIX);
   this.workers =
       EsExecutors.newFixed(
           LOCAL_TRANSPORT_THREAD_NAME_PREFIX,
           workerCount,
           queueSize,
           threadFactory,
           threadPool.getThreadContext());
   this.namedWriteableRegistry = namedWriteableRegistry;
   this.circuitBreakerService = circuitBreakerService;
 }
  protected void messageReceived(
      byte[] data,
      String action,
      LocalTransport sourceTransport,
      Version version,
      @Nullable final Long sendRequestId) {
    Transports.assertTransportThread();
    try {
      transportServiceAdapter.received(data.length);
      StreamInput stream = StreamInput.wrap(data);
      stream.setVersion(version);

      long requestId = stream.readLong();
      byte status = stream.readByte();
      boolean isRequest = TransportStatus.isRequest(status);
      if (isRequest) {
        ThreadContext threadContext = threadPool.getThreadContext();
        threadContext.readHeaders(stream);
        handleRequest(stream, requestId, data.length, sourceTransport, version);
      } else {
        final TransportResponseHandler handler =
            transportServiceAdapter.onResponseReceived(requestId);
        // ignore if its null, the adapter logs it
        if (handler != null) {
          if (TransportStatus.isError(status)) {
            handleResponseError(stream, handler);
          } else {
            handleResponse(stream, sourceTransport, handler);
          }
        }
      }
    } catch (Throwable e) {
      if (sendRequestId != null) {
        TransportResponseHandler handler =
            sourceTransport.transportServiceAdapter.onResponseReceived(sendRequestId);
        if (handler != null) {
          RemoteTransportException error =
              new RemoteTransportException(nodeName(), localAddress, action, e);
          sourceTransport
              .workers()
              .execute(
                  () -> {
                    ThreadContext threadContext = sourceTransport.threadPool.getThreadContext();
                    try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
                      sourceTransport.handleException(handler, error);
                    }
                  });
        }
      } else {
        logger.warn("Failed to receive message for action [{}]", e, action);
      }
    }
  }
 @Override
 protected synchronized void doStart() {
   Objects.requireNonNull(
       clusterStatePublisher, "please set a cluster state publisher before starting");
   Objects.requireNonNull(
       clusterState.nodes().getLocalNode(), "please set the local node before starting");
   Objects.requireNonNull(
       nodeConnectionsService, "please set the node connection service before starting");
   add(localNodeMasterListeners);
   this.clusterState = ClusterState.builder(clusterState).blocks(initialBlocks).build();
   this.updateTasksExecutor =
       EsExecutors.newSinglePrioritizing(
           UPDATE_THREAD_NAME,
           daemonThreadFactory(settings, UPDATE_THREAD_NAME),
           threadPool.getThreadContext());
   this.clusterState = ClusterState.builder(clusterState).blocks(initialBlocks).build();
 }
  @Override
  public void sendRequest(
      final DiscoveryNode node,
      final long requestId,
      final String action,
      final TransportRequest request,
      TransportRequestOptions options)
      throws IOException, TransportException {
    final Version version = Version.smallest(node.getVersion(), this.version);

    try (BytesStreamOutput stream = new BytesStreamOutput()) {
      stream.setVersion(version);

      stream.writeLong(requestId);
      byte status = 0;
      status = TransportStatus.setRequest(status);
      stream.writeByte(status); // 0 for request, 1 for response.

      threadPool.getThreadContext().writeTo(stream);
      stream.writeString(action);
      request.writeTo(stream);

      stream.close();

      final LocalTransport targetTransport = connectedNodes.get(node);
      if (targetTransport == null) {
        throw new NodeNotConnectedException(node, "Node not connected");
      }

      final byte[] data = stream.bytes().toBytes();
      transportServiceAdapter.sent(data.length);
      transportServiceAdapter.onRequestSent(node, requestId, action, request, options);
      targetTransport
          .workers()
          .execute(
              () -> {
                ThreadContext threadContext = targetTransport.threadPool.getThreadContext();
                try (ThreadContext.StoredContext context = threadContext.stashContext()) {
                  targetTransport.messageReceived(
                      data, action, LocalTransport.this, version, requestId);
                }
              });
    }
  }
  private void waitForClusterState(long clusterStateVersion) {
    ClusterStateObserver observer =
        new ClusterStateObserver(
            clusterService, TimeValue.timeValueMinutes(5), logger, threadPool.getThreadContext());
    final ClusterState clusterState = observer.observedState();
    if (clusterState.getVersion() >= clusterStateVersion) {
      logger.trace(
          "node has cluster state with version higher than {} (current: {})",
          clusterStateVersion,
          clusterState.getVersion());
      return;
    } else {
      logger.trace(
          "waiting for cluster state version {} (current: {})",
          clusterStateVersion,
          clusterState.getVersion());
      final PlainActionFuture<Void> future = new PlainActionFuture<>();
      observer.waitForNextChange(
          new ClusterStateObserver.Listener() {

            @Override
            public void onNewClusterState(ClusterState state) {
              future.onResponse(null);
            }

            @Override
            public void onClusterServiceClose() {
              future.onFailure(new NodeClosedException(clusterService.localNode()));
            }

            @Override
            public void onTimeout(TimeValue timeout) {
              future.onFailure(
                  new IllegalStateException(
                      "cluster state never updated to version " + clusterStateVersion));
            }
          },
          new ClusterStateObserver.ValidationPredicate() {

            @Override
            protected boolean validate(ClusterState newState) {
              return newState.getVersion() >= clusterStateVersion;
            }
          });
      try {
        future.get();
        logger.trace(
            "successfully waited for cluster state with version {} (current: {})",
            clusterStateVersion,
            observer.observedState().getVersion());
      } catch (Exception e) {
        logger.debug(
            (Supplier<?>)
                () ->
                    new ParameterizedMessage(
                        "failed waiting for cluster state with version {} (current: {})",
                        clusterStateVersion,
                        observer.observedState()),
            e);
        throw ExceptionsHelper.convertToRuntime(e);
      }
    }
  }
 @Override
 protected boolean apply(String action, ActionRequest request, ActionListener listener) {
   requests.add(new RequestAndHeaders(threadPool.getThreadContext().getHeaders(), request));
   return true;
 }