Skip to content

Commit b3d36d5

Browse files
dnhatnmayya-sharipova
dnhatn
and
mayya-sharipova
authored
Integrate circuit breaker in AsyncTaskIndexService (elastic#73862)
This change integrates the circuit breaker in AsyncTaskIndexService to make sure that we won't hit OOM when serializing a large response of an async search. Related to elastic#67594 Supersedes elastic#73638 Co-authored-by: Mayya Sharipova <[email protected]>
1 parent ff83531 commit b3d36d5

File tree

16 files changed

+231
-62
lines changed

16 files changed

+231
-62
lines changed

‎server/src/main/java/org/elasticsearch/action/update/UpdateRequest.java

+9
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import org.elasticsearch.core.Nullable;
2121
import org.elasticsearch.common.xcontent.ParseField;
2222
import org.elasticsearch.common.Strings;
23+
import org.elasticsearch.common.bytes.BytesReference;
2324
import org.elasticsearch.common.io.stream.StreamInput;
2425
import org.elasticsearch.common.io.stream.StreamOutput;
2526
import org.elasticsearch.common.lucene.uid.Versions;
@@ -659,6 +660,14 @@ public UpdateRequest doc(Object... source) {
659660
return this;
660661
}
661662

663+
/**
664+
* Sets the doc to use for updates when a script is not specified. The doc is provided in a bytes form.
665+
*/
666+
public UpdateRequest doc(BytesReference source, XContentType contentType) {
667+
safeDoc().source(source, contentType);
668+
return this;
669+
}
670+
662671
/**
663672
* Sets the doc to use for updates when a script is not specified, the doc provided
664673
* is a field and value pairs.

‎x-pack/plugin/async-search/src/main/java/org/elasticsearch/xpack/search/TransportGetAsyncSearchAction.java

+7-4
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import org.elasticsearch.cluster.service.ClusterService;
1616
import org.elasticsearch.common.inject.Inject;
1717
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
18+
import org.elasticsearch.common.util.BigArrays;
1819
import org.elasticsearch.tasks.Task;
1920
import org.elasticsearch.threadpool.ThreadPool;
2021
import org.elasticsearch.transport.TransportService;
@@ -37,19 +38,21 @@ public TransportGetAsyncSearchAction(TransportService transportService,
3738
ClusterService clusterService,
3839
NamedWriteableRegistry registry,
3940
Client client,
40-
ThreadPool threadPool) {
41+
ThreadPool threadPool,
42+
BigArrays bigArrays) {
4143
super(GetAsyncSearchAction.NAME, transportService, actionFilters, GetAsyncResultRequest::new);
4244
this.transportService = transportService;
43-
this.resultsService = createResultsService(transportService, clusterService, registry, client, threadPool);
45+
this.resultsService = createResultsService(transportService, clusterService, registry, client, threadPool, bigArrays);
4446
}
4547

4648
static AsyncResultsService<AsyncSearchTask, AsyncSearchResponse> createResultsService(TransportService transportService,
4749
ClusterService clusterService,
4850
NamedWriteableRegistry registry,
4951
Client client,
50-
ThreadPool threadPool) {
52+
ThreadPool threadPool,
53+
BigArrays bigArrays) {
5154
AsyncTaskIndexService<AsyncSearchResponse> store = new AsyncTaskIndexService<>(XPackPlugin.ASYNC_RESULTS_INDEX, clusterService,
52-
threadPool.getThreadContext(), client, ASYNC_SEARCH_ORIGIN, AsyncSearchResponse::new, registry);
55+
threadPool.getThreadContext(), client, ASYNC_SEARCH_ORIGIN, AsyncSearchResponse::new, registry, bigArrays);
5356
return new AsyncResultsService<>(store, true, AsyncSearchTask.class, AsyncSearchTask::addCompletionListener,
5457
transportService.getTaskManager(), clusterService);
5558
}

‎x-pack/plugin/async-search/src/main/java/org/elasticsearch/xpack/search/TransportGetAsyncStatusAction.java

+8-6
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import org.elasticsearch.cluster.service.ClusterService;
1616
import org.elasticsearch.common.inject.Inject;
1717
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
18+
import org.elasticsearch.common.util.BigArrays;
1819
import org.elasticsearch.tasks.Task;
1920
import org.elasticsearch.threadpool.ThreadPool;
2021
import org.elasticsearch.transport.TransportService;
@@ -37,16 +38,17 @@ public class TransportGetAsyncStatusAction extends HandledTransportAction<GetAsy
3738

3839
@Inject
3940
public TransportGetAsyncStatusAction(TransportService transportService,
40-
ActionFilters actionFilters,
41-
ClusterService clusterService,
42-
NamedWriteableRegistry registry,
43-
Client client,
44-
ThreadPool threadPool) {
41+
ActionFilters actionFilters,
42+
ClusterService clusterService,
43+
NamedWriteableRegistry registry,
44+
Client client,
45+
ThreadPool threadPool,
46+
BigArrays bigArrays) {
4547
super(GetAsyncStatusAction.NAME, transportService, actionFilters, GetAsyncStatusRequest::new);
4648
this.transportService = transportService;
4749
this.clusterService = clusterService;
4850
this.store = new AsyncTaskIndexService<>(XPackPlugin.ASYNC_RESULTS_INDEX, clusterService,
49-
threadPool.getThreadContext(), client, ASYNC_SEARCH_ORIGIN, AsyncSearchResponse::new, registry);
51+
threadPool.getThreadContext(), client, ASYNC_SEARCH_ORIGIN, AsyncSearchResponse::new, registry, bigArrays);
5052
}
5153

5254
@Override

‎x-pack/plugin/async-search/src/main/java/org/elasticsearch/xpack/search/TransportSubmitAsyncSearchAction.java

+4-2
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.elasticsearch.common.UUIDs;
2525
import org.elasticsearch.common.inject.Inject;
2626
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
27+
import org.elasticsearch.common.util.BigArrays;
2728
import org.elasticsearch.core.TimeValue;
2829
import org.elasticsearch.common.util.concurrent.ThreadContext;
2930
import org.elasticsearch.index.engine.DocumentMissingException;
@@ -64,14 +65,15 @@ public TransportSubmitAsyncSearchAction(ClusterService clusterService,
6465
Client client,
6566
NodeClient nodeClient,
6667
SearchService searchService,
67-
TransportSearchAction searchAction) {
68+
TransportSearchAction searchAction,
69+
BigArrays bigArrays) {
6870
super(SubmitAsyncSearchAction.NAME, transportService, actionFilters, SubmitAsyncSearchRequest::new);
6971
this.nodeClient = nodeClient;
7072
this.requestToAggReduceContextBuilder = request -> searchService.aggReduceContextBuilder(request).forFinalReduction();
7173
this.searchAction = searchAction;
7274
this.threadContext = transportService.getThreadPool().getThreadContext();
7375
this.store = new AsyncTaskIndexService<>(XPackPlugin.ASYNC_RESULTS_INDEX, clusterService, threadContext, client,
74-
ASYNC_SEARCH_ORIGIN, AsyncSearchResponse::new, registry);
76+
ASYNC_SEARCH_ORIGIN, AsyncSearchResponse::new, registry, bigArrays);
7577
}
7678

7779
@Override

‎x-pack/plugin/async/src/main/java/org/elasticsearch/xpack/async/AsyncResultsIndexPlugin.java

+2-12
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
package org.elasticsearch.xpack.async;
99

1010
import org.elasticsearch.client.Client;
11+
import org.elasticsearch.client.OriginSettingClient;
1112
import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
1213
import org.elasticsearch.cluster.node.DiscoveryNode;
1314
import org.elasticsearch.cluster.service.ClusterService;
@@ -23,10 +24,8 @@
2324
import org.elasticsearch.script.ScriptService;
2425
import org.elasticsearch.threadpool.ThreadPool;
2526
import org.elasticsearch.watcher.ResourceWatcherService;
26-
import org.elasticsearch.xpack.core.XPackPlugin;
2727
import org.elasticsearch.xpack.core.async.AsyncTaskIndexService;
2828
import org.elasticsearch.xpack.core.async.AsyncTaskMaintenanceService;
29-
import org.elasticsearch.xpack.core.search.action.AsyncSearchResponse;
3029

3130
import java.util.ArrayList;
3231
import java.util.Collection;
@@ -75,21 +74,12 @@ public Collection<Object> createComponents(
7574
List<Object> components = new ArrayList<>();
7675
if (DiscoveryNode.canContainData(environment.settings())) {
7776
// only data nodes should be eligible to run the maintenance service.
78-
AsyncTaskIndexService<AsyncSearchResponse> indexService = new AsyncTaskIndexService<>(
79-
XPackPlugin.ASYNC_RESULTS_INDEX,
80-
clusterService,
81-
threadPool.getThreadContext(),
82-
client,
83-
ASYNC_SEARCH_ORIGIN,
84-
AsyncSearchResponse::new,
85-
namedWriteableRegistry
86-
);
8777
AsyncTaskMaintenanceService maintenanceService = new AsyncTaskMaintenanceService(
8878
clusterService,
8979
nodeEnvironment.nodeId(),
9080
settings,
9181
threadPool,
92-
indexService
82+
new OriginSettingClient(client, ASYNC_SEARCH_ORIGIN)
9383
);
9484
components.add(maintenanceService);
9585
}

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/async/AsyncTaskIndexService.java

+26-12
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,11 @@
2626
import org.elasticsearch.common.io.stream.NamedWriteableAwareStreamInput;
2727
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
2828
import org.elasticsearch.common.io.stream.OutputStreamStreamOutput;
29+
import org.elasticsearch.common.io.stream.ReleasableBytesStreamOutput;
2930
import org.elasticsearch.common.io.stream.StreamInput;
3031
import org.elasticsearch.common.io.stream.Writeable;
3132
import org.elasticsearch.common.settings.Settings;
33+
import org.elasticsearch.common.util.BigArrays;
3234
import org.elasticsearch.common.util.concurrent.ThreadContext;
3335
import org.elasticsearch.common.xcontent.XContentBuilder;
3436
import org.elasticsearch.common.xcontent.XContentFactory;
@@ -136,20 +138,23 @@ public static SystemIndexDescriptor getSystemIndexDescriptor() {
136138
private final SecurityContext securityContext;
137139
private final NamedWriteableRegistry registry;
138140
private final Writeable.Reader<R> reader;
141+
private final BigArrays bigArrays;
139142

140143
public AsyncTaskIndexService(String index,
141144
ClusterService clusterService,
142145
ThreadContext threadContext,
143146
Client client,
144147
String origin,
145148
Writeable.Reader<R> reader,
146-
NamedWriteableRegistry registry) {
149+
NamedWriteableRegistry registry,
150+
BigArrays bigArrays) {
147151
this.index = index;
148152
this.securityContext = new SecurityContext(clusterService.getSettings(), threadContext);
149153
this.client = client;
150154
this.clientWithOrigin = new OriginSettingClient(client, origin);
151155
this.registry = registry;
152156
this.reader = reader;
157+
this.bigArrays = bigArrays;
153158
}
154159

155160
/**
@@ -182,18 +187,22 @@ public void createResponse(String docId,
182187
R response,
183188
ActionListener<IndexResponse> listener) throws IOException {
184189
try {
185-
// TODO: Integrate with circuit breaker
186-
final XContentBuilder source = XContentFactory.jsonBuilder()
190+
final ReleasableBytesStreamOutput buffer = new ReleasableBytesStreamOutput(0, bigArrays.withCircuitBreaking());
191+
final XContentBuilder source = XContentFactory.jsonBuilder(buffer);
192+
listener = ActionListener.runBefore(listener, source::close);
193+
source
187194
.startObject()
188195
.field(HEADERS_FIELD, headers)
189196
.field(EXPIRATION_TIME_FIELD, response.getExpirationTime())
190197
.directFieldAsBase64(RESULT_FIELD, os -> writeResponse(response, os))
191198
.endObject();
192-
199+
// do not close the buffer or the XContentBuilder until the IndexRequest is completed (i.e., listener is notified);
200+
// otherwise, we underestimate the memory usage in case the circuit breaker does not use the real memory usage.
201+
source.flush();
193202
final IndexRequest indexRequest = new IndexRequest(index)
194203
.create(true)
195204
.id(docId)
196-
.source(source);
205+
.source(buffer.bytes(), source.contentType());
197206
clientWithOrigin.index(indexRequest, listener);
198207
} catch (Exception e) {
199208
listener.onFailure(e);
@@ -204,20 +213,25 @@ public void createResponse(String docId,
204213
* Stores the final response if the place-holder document is still present (update).
205214
*/
206215
public void updateResponse(String docId,
207-
Map<String, List<String>> responseHeaders,
208-
R response,
209-
ActionListener<UpdateResponse> listener) {
216+
Map<String, List<String>> responseHeaders,
217+
R response,
218+
ActionListener<UpdateResponse> listener) {
210219
try {
211-
// TODO: Integrate with circuit breaker
212-
final XContentBuilder source = XContentFactory.jsonBuilder()
220+
final ReleasableBytesStreamOutput buffer = new ReleasableBytesStreamOutput(0, bigArrays.withCircuitBreaking());
221+
final XContentBuilder source = XContentFactory.jsonBuilder(buffer);
222+
listener = ActionListener.runBefore(listener, source::close);
223+
source
213224
.startObject()
214225
.field(RESPONSE_HEADERS_FIELD, responseHeaders)
215226
.directFieldAsBase64(RESULT_FIELD, os -> writeResponse(response, os))
216227
.endObject();
217-
UpdateRequest request = new UpdateRequest()
228+
// do not close the buffer or the XContentBuilder until the UpdateRequest is completed (i.e., listener is notified);
229+
// otherwise, we underestimate the memory usage in case the circuit breaker does not use the real memory usage.
230+
source.flush();
231+
final UpdateRequest request = new UpdateRequest()
218232
.index(index)
219233
.id(docId)
220-
.doc(source)
234+
.doc(buffer.bytes(), source.contentType())
221235
.retryOnConflict(5);
222236
clientWithOrigin.update(request, listener);
223237
} catch (Exception e) {

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/async/AsyncTaskMaintenanceService.java

+5-5
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.apache.logging.log4j.LogManager;
1111
import org.apache.logging.log4j.Logger;
1212
import org.elasticsearch.action.ActionListener;
13+
import org.elasticsearch.client.Client;
1314
import org.elasticsearch.cluster.ClusterChangedEvent;
1415
import org.elasticsearch.cluster.ClusterState;
1516
import org.elasticsearch.cluster.ClusterStateListener;
@@ -55,7 +56,7 @@ public class AsyncTaskMaintenanceService extends AbstractLifecycleComponent impl
5556
private final String index;
5657
private final String localNodeId;
5758
private final ThreadPool threadPool;
58-
private final AsyncTaskIndexService<?> indexService;
59+
private final Client clientWithOrigin;
5960
private final TimeValue delay;
6061

6162
private boolean isCleanupRunning;
@@ -65,12 +66,12 @@ public AsyncTaskMaintenanceService(ClusterService clusterService,
6566
String localNodeId,
6667
Settings nodeSettings,
6768
ThreadPool threadPool,
68-
AsyncTaskIndexService<?> indexService) {
69+
Client clientWithOrigin) {
6970
this.clusterService = clusterService;
7071
this.index = XPackPlugin.ASYNC_RESULTS_INDEX;
7172
this.localNodeId = localNodeId;
7273
this.threadPool = threadPool;
73-
this.indexService = indexService;
74+
this.clientWithOrigin = clientWithOrigin;
7475
this.delay = ASYNC_SEARCH_CLEANUP_INTERVAL_SETTING.get(nodeSettings);
7576
}
7677

@@ -125,8 +126,7 @@ synchronized void executeNextCleanup() {
125126
long nowInMillis = System.currentTimeMillis();
126127
DeleteByQueryRequest toDelete = new DeleteByQueryRequest(index)
127128
.setQuery(QueryBuilders.rangeQuery(EXPIRATION_TIME_FIELD).lte(nowInMillis));
128-
indexService.getClientWithOrigin()
129-
.execute(DeleteByQueryAction.INSTANCE, toDelete, ActionListener.wrap(this::scheduleNextCleanup));
129+
clientWithOrigin.execute(DeleteByQueryAction.INSTANCE, toDelete, ActionListener.wrap(this::scheduleNextCleanup));
130130
}
131131
}
132132

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/async/TransportDeleteAsyncResultAction.java

+4-2
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import org.elasticsearch.cluster.service.ClusterService;
1717
import org.elasticsearch.common.inject.Inject;
1818
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
19+
import org.elasticsearch.common.util.BigArrays;
1920
import org.elasticsearch.tasks.Task;
2021
import org.elasticsearch.threadpool.ThreadPool;
2122
import org.elasticsearch.transport.TransportService;
@@ -34,13 +35,14 @@ public TransportDeleteAsyncResultAction(TransportService transportService,
3435
ClusterService clusterService,
3536
NamedWriteableRegistry registry,
3637
Client client,
37-
ThreadPool threadPool) {
38+
ThreadPool threadPool,
39+
BigArrays bigArrays) {
3840
super(DeleteAsyncResultAction.NAME, transportService, actionFilters, DeleteAsyncResultRequest::new);
3941
this.transportService = transportService;
4042
this.clusterService = clusterService;
4143
AsyncTaskIndexService<?> store = new AsyncTaskIndexService<>(XPackPlugin.ASYNC_RESULTS_INDEX, clusterService,
4244
threadPool.getThreadContext(), client, ASYNC_SEARCH_ORIGIN,
43-
(in) -> {throw new UnsupportedOperationException("Reading is not supported during deletion");}, registry);
45+
(in) -> {throw new UnsupportedOperationException("Reading is not supported during deletion");}, registry, bigArrays);
4446
this.deleteResultsService = new DeleteAsyncResultsService(store, transportService.getTaskManager());
4547
}
4648

‎x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/async/AsyncResultsServiceTests.java

+3-1
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.elasticsearch.action.update.UpdateResponse;
1515
import org.elasticsearch.cluster.service.ClusterService;
1616
import org.elasticsearch.core.TimeValue;
17+
import org.elasticsearch.common.util.BigArrays;
1718
import org.elasticsearch.tasks.CancellableTask;
1819
import org.elasticsearch.tasks.Task;
1920
import org.elasticsearch.tasks.TaskId;
@@ -123,9 +124,10 @@ public Task createTask(long id, String type, String action, TaskId parentTaskId,
123124
public void setup() {
124125
clusterService = getInstanceFromNode(ClusterService.class);
125126
TransportService transportService = getInstanceFromNode(TransportService.class);
127+
BigArrays bigArrays = getInstanceFromNode(BigArrays.class);
126128
taskManager = transportService.getTaskManager();
127129
indexService = new AsyncTaskIndexService<>("test", clusterService, transportService.getThreadPool().getThreadContext(),
128-
client(), ASYNC_SEARCH_ORIGIN, TestAsyncResponse::new, writableRegistry());
130+
client(), ASYNC_SEARCH_ORIGIN, TestAsyncResponse::new, writableRegistry(), bigArrays);
129131

130132
}
131133

0 commit comments

Comments
 (0)