-
Notifications
You must be signed in to change notification settings - Fork 24.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Browse files
Browse the repository at this point in the history
We have implemented #99927 in DriverRunner. However, we also need to implement this in ComputeService, where we spawn multiple requests to avoid losing response headers. Relates #99927 Closes #100163 Closes #102982 Closes #102871 Closes #103028
- Loading branch information
Showing
6 changed files
with
253 additions
and
33 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,9 @@ | ||
| pr: 103031 | ||
| summary: Collect warnings in compute service | ||
| area: ES|QL | ||
| type: bug | ||
| issues: | ||
| - 100163 | ||
| - 103028 | ||
| - 102871 | ||
| - 102982 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
60 changes: 60 additions & 0 deletions
60
...ql/compute/src/main/java/org/elasticsearch/compute/operator/ResponseHeadersCollector.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,60 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.compute.operator; | ||
|
|
||
| import org.elasticsearch.common.util.concurrent.ConcurrentCollections; | ||
| import org.elasticsearch.common.util.concurrent.ThreadContext; | ||
|
|
||
| import java.util.HashMap; | ||
| import java.util.LinkedHashSet; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.Queue; | ||
| import java.util.Set; | ||
|
|
||
| /** | ||
| * A helper class that can be used to collect and merge response headers from multiple child requests. | ||
| */ | ||
| public final class ResponseHeadersCollector { | ||
| private final ThreadContext threadContext; | ||
| private final Queue<Map<String, List<String>>> collected = ConcurrentCollections.newQueue(); | ||
|
|
||
| public ResponseHeadersCollector(ThreadContext threadContext) { | ||
| this.threadContext = threadContext; | ||
| } | ||
|
|
||
| /** | ||
| * Called when a child request is completed to collect the response headers of the responding thread | ||
| */ | ||
| public void collect() { | ||
| Map<String, List<String>> responseHeaders = threadContext.getResponseHeaders(); | ||
| if (responseHeaders.isEmpty() == false) { | ||
| collected.add(responseHeaders); | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Called when all child requests are completed. This will merge all collected response headers | ||
| * from the child requests and restore to the current thread. | ||
| */ | ||
| public void finish() { | ||
| final Map<String, Set<String>> merged = new HashMap<>(); | ||
| Map<String, List<String>> resp; | ||
| while ((resp = collected.poll()) != null) { | ||
| for (Map.Entry<String, List<String>> e : resp.entrySet()) { | ||
| // Use LinkedHashSet to retain the order of the values | ||
| merged.computeIfAbsent(e.getKey(), k -> new LinkedHashSet<>(e.getValue().size())).addAll(e.getValue()); | ||
| } | ||
| } | ||
| for (Map.Entry<String, Set<String>> e : merged.entrySet()) { | ||
| for (String v : e.getValue()) { | ||
| threadContext.addResponseHeader(e.getKey(), v); | ||
| } | ||
| } | ||
| } | ||
| } |
72 changes: 72 additions & 0 deletions
72
...mpute/src/test/java/org/elasticsearch/compute/operator/ResponseHeadersCollectorTests.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,72 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.compute.operator; | ||
|
|
||
| import org.elasticsearch.action.ActionListener; | ||
| import org.elasticsearch.action.ActionRunnable; | ||
| import org.elasticsearch.action.support.PlainActionFuture; | ||
| import org.elasticsearch.action.support.RefCountingListener; | ||
| import org.elasticsearch.common.settings.Settings; | ||
| import org.elasticsearch.common.util.concurrent.EsExecutors; | ||
| import org.elasticsearch.common.util.concurrent.ThreadContext; | ||
| import org.elasticsearch.common.util.set.Sets; | ||
| import org.elasticsearch.core.TimeValue; | ||
| import org.elasticsearch.test.ESTestCase; | ||
| import org.elasticsearch.threadpool.FixedExecutorBuilder; | ||
| import org.elasticsearch.threadpool.TestThreadPool; | ||
|
|
||
| import java.util.HashSet; | ||
| import java.util.List; | ||
| import java.util.Set; | ||
| import java.util.concurrent.CyclicBarrier; | ||
| import java.util.concurrent.TimeUnit; | ||
|
|
||
| import static org.hamcrest.Matchers.equalTo; | ||
|
|
||
| public class ResponseHeadersCollectorTests extends ESTestCase { | ||
|
|
||
| public void testCollect() { | ||
| int numThreads = randomIntBetween(1, 10); | ||
| TestThreadPool threadPool = new TestThreadPool( | ||
| getTestClass().getSimpleName(), | ||
| new FixedExecutorBuilder(Settings.EMPTY, "test", numThreads, 1024, "test", EsExecutors.TaskTrackingConfig.DEFAULT) | ||
| ); | ||
| Set<String> expectedWarnings = new HashSet<>(); | ||
| try { | ||
| ThreadContext threadContext = threadPool.getThreadContext(); | ||
| var collector = new ResponseHeadersCollector(threadContext); | ||
| PlainActionFuture<Void> future = new PlainActionFuture<>(); | ||
| Runnable mergeAndVerify = () -> { | ||
| collector.finish(); | ||
| List<String> actualWarnings = threadContext.getResponseHeaders().getOrDefault("Warnings", List.of()); | ||
| assertThat(Sets.newHashSet(actualWarnings), equalTo(expectedWarnings)); | ||
| }; | ||
| try (RefCountingListener refs = new RefCountingListener(ActionListener.runAfter(future, mergeAndVerify))) { | ||
| CyclicBarrier barrier = new CyclicBarrier(numThreads); | ||
| for (int i = 0; i < numThreads; i++) { | ||
| String warning = "warning-" + i; | ||
| expectedWarnings.add(warning); | ||
| ActionListener<Void> listener = ActionListener.runBefore(refs.acquire(), collector::collect); | ||
| threadPool.schedule(new ActionRunnable<>(listener) { | ||
| @Override | ||
| protected void doRun() throws Exception { | ||
| barrier.await(30, TimeUnit.SECONDS); | ||
| try (ThreadContext.StoredContext ignored = threadContext.stashContext()) { | ||
| threadContext.addResponseHeader("Warnings", warning); | ||
| listener.onResponse(null); | ||
| } | ||
| } | ||
| }, TimeValue.timeValueNanos(between(0, 1000_000)), threadPool.executor("test")); | ||
| } | ||
| } | ||
| future.actionGet(TimeValue.timeValueSeconds(30)); | ||
| } finally { | ||
| terminate(threadPool); | ||
| } | ||
| } | ||
| } |
85 changes: 85 additions & 0 deletions
85
...gin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/WarningsIT.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,85 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.esql.action; | ||
|
|
||
| import org.elasticsearch.action.ActionListener; | ||
| import org.elasticsearch.action.support.PlainActionFuture; | ||
| import org.elasticsearch.cluster.node.DiscoveryNode; | ||
| import org.elasticsearch.common.settings.Settings; | ||
| import org.elasticsearch.test.junit.annotations.TestLogging; | ||
| import org.elasticsearch.transport.TransportService; | ||
|
|
||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.concurrent.TimeUnit; | ||
|
|
||
| import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked; | ||
| import static org.hamcrest.Matchers.greaterThanOrEqualTo; | ||
|
|
||
| @TestLogging(value = "org.elasticsearch.xpack.esql:TRACE", reason = "debug") | ||
| public class WarningsIT extends AbstractEsqlIntegTestCase { | ||
|
|
||
| public void testCollectWarnings() { | ||
| final String node1, node2; | ||
| if (randomBoolean()) { | ||
| internalCluster().ensureAtLeastNumDataNodes(2); | ||
| node1 = randomDataNode().getName(); | ||
| node2 = randomValueOtherThan(node1, () -> randomDataNode().getName()); | ||
| } else { | ||
| node1 = randomDataNode().getName(); | ||
| node2 = randomDataNode().getName(); | ||
| } | ||
|
|
||
| int numDocs1 = randomIntBetween(1, 15); | ||
| assertAcked( | ||
| client().admin() | ||
| .indices() | ||
| .prepareCreate("index-1") | ||
| .setSettings(Settings.builder().put("index.routing.allocation.require._name", node1)) | ||
| .setMapping("host", "type=keyword") | ||
| ); | ||
| for (int i = 0; i < numDocs1; i++) { | ||
| client().prepareIndex("index-1").setSource("host", "192." + i).get(); | ||
| } | ||
| int numDocs2 = randomIntBetween(1, 15); | ||
| assertAcked( | ||
| client().admin() | ||
| .indices() | ||
| .prepareCreate("index-2") | ||
| .setSettings(Settings.builder().put("index.routing.allocation.require._name", node2)) | ||
| .setMapping("host", "type=keyword") | ||
| ); | ||
| for (int i = 0; i < numDocs2; i++) { | ||
| client().prepareIndex("index-2").setSource("host", "10." + i).get(); | ||
| } | ||
|
|
||
| DiscoveryNode coordinator = randomFrom(clusterService().state().nodes().stream().toList()); | ||
| client().admin().indices().prepareRefresh("index-1", "index-2").get(); | ||
|
|
||
| EsqlQueryRequest request = new EsqlQueryRequest(); | ||
| request.query("FROM index-* | EVAL ip = to_ip(host) | STATS s = COUNT(*) by ip | KEEP ip | LIMIT 100"); | ||
| request.pragmas(randomPragmas()); | ||
| PlainActionFuture<EsqlQueryResponse> future = new PlainActionFuture<>(); | ||
| client(coordinator.getName()).execute(EsqlQueryAction.INSTANCE, request, ActionListener.runBefore(future, () -> { | ||
| var threadpool = internalCluster().getInstance(TransportService.class, coordinator.getName()).getThreadPool(); | ||
| Map<String, List<String>> responseHeaders = threadpool.getThreadContext().getResponseHeaders(); | ||
| List<String> warnings = responseHeaders.getOrDefault("Warning", List.of()) | ||
| .stream() | ||
| .filter(w -> w.contains("is not an IP string literal")) | ||
| .toList(); | ||
| int expectedWarnings = Math.min(20, numDocs1 + numDocs2); | ||
| // we cap the number of warnings per node | ||
| assertThat(warnings.size(), greaterThanOrEqualTo(expectedWarnings)); | ||
| })); | ||
| future.actionGet(30, TimeUnit.SECONDS).close(); | ||
| } | ||
|
|
||
| private DiscoveryNode randomDataNode() { | ||
| return randomFrom(clusterService().state().nodes().getDataNodes().values()); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters