|
@@ -8,6 +8,7 @@
|
|
|
package org.elasticsearch.xpack.esql.action;
|
|
|
|
|
|
import org.elasticsearch.ElasticsearchException;
|
|
|
+import org.elasticsearch.ExceptionsHelper;
|
|
|
import org.elasticsearch.action.DocWriteResponse;
|
|
|
import org.elasticsearch.cluster.metadata.IndexMetadata;
|
|
|
import org.elasticsearch.common.Strings;
|
|
@@ -20,6 +21,7 @@ import org.elasticsearch.compute.operator.exchange.ExchangeService;
|
|
|
import org.elasticsearch.core.TimeValue;
|
|
|
import org.elasticsearch.indices.breaker.HierarchyCircuitBreakerService;
|
|
|
import org.elasticsearch.plugins.Plugin;
|
|
|
+import org.elasticsearch.rest.RestStatus;
|
|
|
import org.elasticsearch.test.junit.annotations.TestLogging;
|
|
|
|
|
|
import java.util.ArrayList;
|
|
@@ -76,18 +78,31 @@ public class EsqlActionBreakerIT extends EsqlActionIT {
|
|
|
.build();
|
|
|
}
|
|
|
|
|
|
- @Override
|
|
|
- protected EsqlQueryResponse run(EsqlQueryRequest request) {
|
|
|
+ private EsqlQueryResponse runWithBreaking(EsqlQueryRequest request) throws CircuitBreakingException {
|
|
|
setRequestCircuitBreakerLimit(ByteSizeValue.ofBytes(between(256, 2048)));
|
|
|
try {
|
|
|
return client().execute(EsqlQueryAction.INSTANCE, request).actionGet(2, TimeUnit.MINUTES);
|
|
|
} catch (Exception e) {
|
|
|
logger.info("request failed", e);
|
|
|
ensureBlocksReleased();
|
|
|
+ throw e;
|
|
|
} finally {
|
|
|
setRequestCircuitBreakerLimit(null);
|
|
|
}
|
|
|
- return super.run(request);
|
|
|
+ }
|
|
|
+
|
|
|
+ @Override
|
|
|
+ protected EsqlQueryResponse run(EsqlQueryRequest request) {
|
|
|
+ try {
|
|
|
+ return runWithBreaking(request);
|
|
|
+ } catch (Exception e) {
|
|
|
+ try (EsqlQueryResponse resp = super.run(request)) {
|
|
|
+ assertThat(e, instanceOf(CircuitBreakingException.class));
|
|
|
+ assertThat(ExceptionsHelper.status(e), equalTo(RestStatus.TOO_MANY_REQUESTS));
|
|
|
+ resp.incRef();
|
|
|
+ return resp;
|
|
|
+ }
|
|
|
+ }
|
|
|
}
|
|
|
|
|
|
/**
|