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;
}