|
11 | 11 | import org.elasticsearch.action.ActionListenerResponseHandler; |
12 | 12 | import org.elasticsearch.action.OriginalIndices; |
13 | 13 | import org.elasticsearch.action.support.ChannelActionListener; |
| 14 | +import org.elasticsearch.compute.operator.DriverProfile; |
14 | 15 | import org.elasticsearch.compute.operator.exchange.ExchangeService; |
15 | 16 | import org.elasticsearch.compute.operator.exchange.ExchangeSourceHandler; |
16 | 17 | import org.elasticsearch.core.Releasable; |
|
34 | 35 | import java.util.ArrayList; |
35 | 36 | import java.util.List; |
36 | 37 | import java.util.Map; |
37 | | -import java.util.Objects; |
38 | 38 | import java.util.Set; |
39 | 39 | import java.util.concurrent.Executor; |
| 40 | +import java.util.concurrent.atomic.AtomicBoolean; |
40 | 41 | import java.util.concurrent.atomic.AtomicReference; |
41 | 42 |
|
42 | 43 | /** |
@@ -74,69 +75,97 @@ void startComputeOnRemoteCluster( |
74 | 75 | RemoteCluster cluster, |
75 | 76 | Runnable cancelQueryOnFailure, |
76 | 77 | EsqlExecutionInfo executionInfo, |
77 | | - ActionListener<ComputeResponse> listener |
| 78 | + ActionListener<List<DriverProfile>> listener |
78 | 79 | ) { |
79 | 80 | var queryPragmas = configuration.pragmas(); |
80 | 81 | listener = ActionListener.runBefore(listener, exchangeSource.addEmptySink()::close); |
81 | 82 | final var childSessionId = computeService.newChildSession(sessionId); |
82 | | - final AtomicReference<ComputeResponse> finalResponse = new AtomicReference<>(); |
83 | 83 | final String clusterAlias = cluster.clusterAlias(); |
84 | | - try (var computeListener = new ComputeListener(transportService.getThreadPool(), cancelQueryOnFailure, listener.map(profiles -> { |
85 | | - var resp = finalResponse.get(); |
86 | | - return Objects.requireNonNullElseGet(resp, () -> new ComputeResponse(profiles)); |
87 | | - }))) { |
88 | | - ExchangeService.openExchange( |
89 | | - transportService, |
90 | | - cluster.connection, |
91 | | - childSessionId, |
92 | | - queryPragmas.exchangeBufferSize(), |
93 | | - esqlExecutor, |
94 | | - EsqlCCSUtils.skipUnavailableListener( |
95 | | - computeListener.acquireAvoid(), |
96 | | - executionInfo, |
97 | | - clusterAlias, |
98 | | - EsqlExecutionInfo.Cluster.Status.SKIPPED |
99 | | - ).delegateFailureAndWrap((l, unused) -> { |
100 | | - var listenerGroup = new RemoteListenerGroup( |
101 | | - transportService, |
102 | | - rootTask, |
103 | | - computeListener, |
104 | | - clusterAlias, |
105 | | - executionInfo, |
106 | | - l |
107 | | - ); |
108 | | - |
109 | | - var remoteSink = exchangeService.newRemoteSink( |
110 | | - listenerGroup.getGroupTask(), |
111 | | - childSessionId, |
112 | | - transportService, |
113 | | - cluster.connection |
114 | | - ); |
| 84 | + final AtomicBoolean pagesFetched = new AtomicBoolean(); |
| 85 | + final AtomicReference<ComputeResponse> finalResponse = new AtomicReference<>(); |
| 86 | + listener = listener.delegateResponse((l, e) -> { |
| 87 | + final boolean receivedResults = finalResponse.get() != null || pagesFetched.get(); |
| 88 | + if (receivedResults == false && EsqlCCSUtils.shouldIgnoreRuntimeError(executionInfo, clusterAlias, e)) { |
| 89 | + EsqlCCSUtils.markClusterWithFinalStateAndNoShards(executionInfo, clusterAlias, EsqlExecutionInfo.Cluster.Status.SKIPPED, e); |
| 90 | + l.onResponse(List.of()); |
| 91 | + } else if (configuration.allowPartialResults()) { |
| 92 | + EsqlCCSUtils.markClusterWithFinalStateAndNoShards(executionInfo, clusterAlias, EsqlExecutionInfo.Cluster.Status.PARTIAL, e); |
| 93 | + l.onResponse(List.of()); |
| 94 | + } else { |
| 95 | + l.onFailure(e); |
| 96 | + } |
| 97 | + }); |
| 98 | + ExchangeService.openExchange( |
| 99 | + transportService, |
| 100 | + cluster.connection, |
| 101 | + childSessionId, |
| 102 | + queryPragmas.exchangeBufferSize(), |
| 103 | + esqlExecutor, |
| 104 | + listener.delegateFailure((l, unused) -> { |
| 105 | + final CancellableTask groupTask; |
| 106 | + final Runnable onGroupFailure; |
| 107 | + boolean failFast = executionInfo.isSkipUnavailable(clusterAlias) == false && configuration.allowPartialResults() == false; |
| 108 | + if (failFast) { |
| 109 | + groupTask = rootTask; |
| 110 | + onGroupFailure = cancelQueryOnFailure; |
| 111 | + } else { |
| 112 | + groupTask = computeService.createGroupTask(rootTask, () -> "compute group: cluster [" + clusterAlias + "]"); |
| 113 | + onGroupFailure = computeService.cancelQueryOnFailure(groupTask); |
| 114 | + l = ActionListener.runAfter(l, () -> transportService.getTaskManager().unregister(groupTask)); |
| 115 | + } |
| 116 | + try (var computeListener = new ComputeListener(transportService.getThreadPool(), onGroupFailure, l.map(profiles -> { |
| 117 | + updateExecutionInfo(executionInfo, clusterAlias, finalResponse.get()); |
| 118 | + return profiles; |
| 119 | + }))) { |
| 120 | + var remoteSink = exchangeService.newRemoteSink(groupTask, childSessionId, transportService, cluster.connection); |
115 | 121 | exchangeSource.addRemoteSink( |
116 | 122 | remoteSink, |
117 | | - executionInfo.isSkipUnavailable(clusterAlias) == false, |
118 | | - () -> {}, |
| 123 | + failFast, |
| 124 | + () -> pagesFetched.set(true), |
119 | 125 | queryPragmas.concurrentExchangeClients(), |
120 | | - listenerGroup.getExchangeRequestListener() |
| 126 | + computeListener.acquireAvoid() |
121 | 127 | ); |
122 | 128 | var remotePlan = new RemoteClusterPlan(plan, cluster.concreteIndices, cluster.originalIndices); |
123 | 129 | var clusterRequest = new ClusterComputeRequest(clusterAlias, childSessionId, configuration, remotePlan); |
124 | | - final ActionListener<ComputeResponse> clusterListener = listenerGroup.getClusterRequestListener().map(r -> { |
| 130 | + final ActionListener<ComputeResponse> clusterListener = computeListener.acquireCompute().map(r -> { |
125 | 131 | finalResponse.set(r); |
126 | 132 | return r.getProfiles(); |
127 | 133 | }); |
128 | 134 | transportService.sendChildRequest( |
129 | 135 | cluster.connection, |
130 | 136 | ComputeService.CLUSTER_ACTION_NAME, |
131 | 137 | clusterRequest, |
132 | | - listenerGroup.getGroupTask(), |
| 138 | + groupTask, |
133 | 139 | TransportRequestOptions.EMPTY, |
134 | 140 | new ActionListenerResponseHandler<>(clusterListener, ComputeResponse::new, esqlExecutor) |
135 | 141 | ); |
136 | | - }) |
137 | | - ); |
138 | | - } |
| 142 | + } |
| 143 | + }) |
| 144 | + ); |
| 145 | + } |
139 | 146 |
|
| 147 | + private void updateExecutionInfo(EsqlExecutionInfo executionInfo, String clusterAlias, ComputeResponse resp) { |
| 148 | + executionInfo.swapCluster(clusterAlias, (k, v) -> { |
| 149 | + var builder = new EsqlExecutionInfo.Cluster.Builder(v).setTotalShards(resp.getTotalShards()) |
| 150 | + .setSuccessfulShards(resp.getSuccessfulShards()) |
| 151 | + .setSkippedShards(resp.getSkippedShards()) |
| 152 | + .setFailedShards(resp.getFailedShards()); |
| 153 | + if (resp.getTook() != null) { |
| 154 | + builder.setTook(TimeValue.timeValueNanos(executionInfo.planningTookTime().nanos() + resp.getTook().nanos())); |
| 155 | + } else { |
| 156 | + // if the cluster is an older version and does not send back took time, then calculate it here on the coordinator |
| 157 | + // and leave shard info unset, so it is not shown in the CCS metadata section of the JSON response |
| 158 | + builder.setTook(executionInfo.tookSoFar()); |
| 159 | + } |
| 160 | + if (v.getStatus() == EsqlExecutionInfo.Cluster.Status.RUNNING) { |
| 161 | + if (executionInfo.isStopped() || resp.failedShards > 0) { |
| 162 | + builder.setStatus(EsqlExecutionInfo.Cluster.Status.PARTIAL); |
| 163 | + } else { |
| 164 | + builder.setStatus(EsqlExecutionInfo.Cluster.Status.SUCCESSFUL); |
| 165 | + } |
| 166 | + } |
| 167 | + return builder.build(); |
| 168 | + }); |
140 | 169 | } |
141 | 170 |
|
142 | 171 | List<RemoteCluster> getRemoteClusters( |
|
0 commit comments