Skip to content

Make ESQL EnrichPolicyResolver try to do proper connection before sending requests #114870

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Oct 16, 2024
Merged

Make ESQL EnrichPolicyResolver try to do proper connection before sending requests #114870

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
17428eb
Make ESQL EnrichPolicyResolver try to do proper connection before sen…
smalyshev Oct 15, 2024
12ebeb2
Make encureConnected be !skipUnavailable
smalyshev Oct 16, 2024
44df371
Merge branch 'main' into esql-reconnect
smalyshev Oct 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 25 additions & 20 deletions ...k/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichPolicyResolver.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import org.elasticsearch.tasks.Task;
import org.elasticsearch.threadpool.ThreadPool;
import org.elasticsearch.transport.RemoteClusterAware;
import org.elasticsearch.transport.RemoteClusterService;
import org.elasticsearch.transport.Transport;
import org.elasticsearch.transport.TransportChannel;
import org.elasticsearch.transport.TransportRequest;
Expand Down Expand Up @@ -72,12 +73,14 @@ public class EnrichPolicyResolver {
private final IndexResolver indexResolver;
private final TransportService transportService;
private final ThreadPool threadPool;
private final RemoteClusterService remoteClusterService;

public EnrichPolicyResolver(ClusterService clusterService, TransportService transportService, IndexResolver indexResolver) {
this.clusterService = clusterService;
this.transportService = transportService;
this.indexResolver = indexResolver;
this.threadPool = transportService.getThreadPool();
this.remoteClusterService = transportService.getRemoteClusterService();
transportService.registerRequestHandler(
RESOLVE_ACTION_NAME,
threadPool.executor(ThreadPool.Names.SEARCH),
Expand Down Expand Up @@ -257,22 +260,21 @@ private void lookupPolicies(
// remote clusters
if (remotePolicies.isEmpty() == false) {
for (String cluster : remoteClusters) {
final Transport.Connection connection;
try {
connection = getRemoteConnection(cluster);
} catch (Exception e) {
refs.acquire().onFailure(e);
return;
}
transportService.sendRequest(
connection,
RESOLVE_ACTION_NAME,
new LookupRequest(cluster, remotePolicies),
TransportRequestOptions.EMPTY,
new ActionListenerResponseHandler<>(
refs.acquire(resp -> lookupResponses.put(cluster, resp)),
LookupResponse::new,
threadPool.executor(ThreadPool.Names.SEARCH)
ActionListener<LookupResponse> lookupListener = refs.acquire(resp -> lookupResponses.put(cluster, resp));
getRemoteConnection(
cluster,
lookupListener.delegateFailureAndWrap(
(delegate, connection) -> transportService.sendRequest(
connection,
RESOLVE_ACTION_NAME,
new LookupRequest(cluster, remotePolicies),
TransportRequestOptions.EMPTY,
new ActionListenerResponseHandler<>(
delegate,
LookupResponse::new,
threadPool.executor(ThreadPool.Names.SEARCH)
)
)
)
);
}
Expand Down Expand Up @@ -389,13 +391,16 @@ protected Map<String, EnrichPolicy> availablePolicies() {
return metadata == null ? Map.of() : metadata.getPolicies();
}

protected Transport.Connection getRemoteConnection(String cluster) {
return transportService.getRemoteClusterService().getConnection(cluster);
protected void getRemoteConnection(String cluster, ActionListener<Transport.Connection> listener) {
remoteClusterService.maybeEnsureConnectedAndGetConnection(
cluster,
remoteClusterService.isSkipUnavailable(cluster) == false,
listener
);
}

public Map<String, List<String>> groupIndicesPerCluster(String[] indices) {
return transportService.getRemoteClusterService()
.groupIndices(SearchRequest.DEFAULT_INDICES_OPTIONS, indices)
return remoteClusterService.groupIndices(SearchRequest.DEFAULT_INDICES_OPTIONS, indices)
.entrySet()
.stream()
.collect(Collectors.toMap(Map.Entry::getKey, e -> Arrays.asList(e.getValue().indices())));
Expand Down
4 changes: 2 additions & 2 deletions ...gin/esql/src/test/java/org/elasticsearch/xpack/esql/enrich/EnrichPolicyResolverTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -446,9 +446,9 @@ EnrichResolution resolvePolicies(Collection<String> clusters, Collection<Unresol
}

@Override
protected Transport.Connection getRemoteConnection(String remoteCluster) {
protected void getRemoteConnection(String remoteCluster, ActionListener<Transport.Connection> listener) {
assertThat("Must only called on the local cluster", cluster, equalTo(LOCAL_CLUSTER_GROUP_KEY));
return transports.get("").getConnection(transports.get(remoteCluster).getLocalNode());
listener.onResponse(transports.get("").getConnection(transports.get(remoteCluster).getLocalNode()));
}

static ClusterService mockClusterService(Map<String, EnrichPolicy> policies) {
Expand Down