add HeapBufferedAsyncResponseConsumer with configurable max buffer size

We keep the default async client behaviour like in BasicAsyncResponseConsumer, but we lower the maximum size of the buffer from Integer.MAX_VALUE (2GB) to 10 MB. This way users will realize they are buffering big responses in heap hence they'll know they have to do something about it, either write their own response consumer or increase the buffer size limit by providing their manually creeted instance of HeapBufferedAsyncResponseConsumer (constructor accept a bufferLimit int argument).
This commit is contained in:
javanna 2016-07-19 14:51:46 +02:00 committed by Luca Cavanna
parent 1bb33cf572
commit 8eccdff9ad
4 changed files with 243 additions and 5 deletions

View File

@ -0,0 +1,108 @@
/*
* Licensed to Elasticsearch under one or more contributor
* license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright
* ownership. Elasticsearch licenses this file to you under
* the Apache License, Version 2.0 (the "License"); you may
* not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.elasticsearch.client;
import org.apache.http.ContentTooLongException;
import org.apache.http.HttpEntity;
import org.apache.http.HttpException;
import org.apache.http.HttpResponse;
import org.apache.http.entity.ContentType;
import org.apache.http.nio.ContentDecoder;
import org.apache.http.nio.IOControl;
import org.apache.http.nio.entity.ContentBufferEntity;
import org.apache.http.nio.protocol.AbstractAsyncResponseConsumer;
import org.apache.http.nio.util.ByteBufferAllocator;
import org.apache.http.nio.util.HeapByteBufferAllocator;
import org.apache.http.nio.util.SimpleInputBuffer;
import org.apache.http.protocol.HttpContext;
import java.io.IOException;
/**
* Default implementation of {@link org.apache.http.nio.protocol.HttpAsyncResponseConsumer}. Buffers the whole
* response content in heap memory, meaning that the size of the buffer is equal to the content-length of the response.
* Limits the size of responses that can be read to {@link #DEFAULT_BUFFER_LIMIT} by default, configurable value.
* Throws an exception in case the entity is longer than the configured buffer limit.
*/
public class HeapBufferedAsyncResponseConsumer extends AbstractAsyncResponseConsumer<HttpResponse> {
//default buffer limit is 10MB
public static final int DEFAULT_BUFFER_LIMIT = 10 * 1024 * 1024;
private final int bufferLimit;
private volatile HttpResponse response;
private volatile SimpleInputBuffer buf;
/**
* Creates a new instance of this consumer with a buffer limit of {@link #DEFAULT_BUFFER_LIMIT}
*/
public HeapBufferedAsyncResponseConsumer() {
this.bufferLimit = DEFAULT_BUFFER_LIMIT;
}
/**
* Creates a new instance of this consumer with the provided buffer limit
*/
public HeapBufferedAsyncResponseConsumer(int bufferLimit) {
this.bufferLimit = bufferLimit;
}
@Override
protected void onResponseReceived(HttpResponse response) throws HttpException, IOException {
this.response = response;
}
@Override
protected void onEntityEnclosed(HttpEntity entity, ContentType contentType) throws IOException {
long len = entity.getContentLength();
if (len > bufferLimit) {
throw new ContentTooLongException("entity content is too long [" + len +
"] for the configured buffer limit [" + bufferLimit + "]");
}
if (len < 0) {
len = 4096;
}
this.buf = new SimpleInputBuffer((int) len, getByteBufferAllocator());
this.response.setEntity(new ContentBufferEntity(entity, this.buf));
}
/**
* Returns the instance of {@link ByteBufferAllocator} to use for content buffering.
* Allows to plug in any {@link ByteBufferAllocator} implementation.
*/
protected ByteBufferAllocator getByteBufferAllocator() {
return HeapByteBufferAllocator.INSTANCE;
}
@Override
protected void onContentReceived(ContentDecoder decoder, IOControl ioctrl) throws IOException {
this.buf.consumeContent(decoder);
}
@Override
protected HttpResponse buildResult(HttpContext context) throws Exception {
return response;
}
@Override
protected void releaseResources() {
response = null;
}
}

View File

@ -25,7 +25,7 @@ import java.io.IOException;
* Exception thrown when an elasticsearch node responds to a request with a status code that indicates an error.
* Holds the response that was returned.
*/
public class ResponseException extends IOException {
public final class ResponseException extends IOException {
private Response response;

View File

@ -172,9 +172,34 @@ public final class RestClient implements Closeable {
*/
public Response performRequest(String method, String endpoint, Map<String, String> params,
HttpEntity entity, Header... headers) throws Exception {
HttpAsyncResponseConsumer<HttpResponse> consumer = HttpAsyncMethods.createConsumer();
HttpAsyncResponseConsumer<HttpResponse> responseConsumer = new HeapBufferedAsyncResponseConsumer();
return performRequest(method, endpoint, params, entity, responseConsumer, headers);
}
/**
* Sends a request to the elasticsearch cluster that the current client points to. Blocks until the request is completed and returns
* its response of fails by throwing an exception. Selects a host out of the provided ones in a round-robin fashion. Failing hosts
* are marked dead and retried after a certain amount of time (minimum 1 minute, maximum 30 minutes), depending on how many times
* they previously failed (the more failures, the later they will be retried). In case of failures all of the alive nodes (or dead
* nodes that deserve a retry) are retried until one responds or none of them does, in which case an {@link IOException} will be thrown.
*
* @param method the http method
* @param endpoint the path of the request (without host and port)
* @param params the query_string parameters
* @param entity the body of the request, null if not applicable
* @param responseConsumer the {@link HttpAsyncResponseConsumer} callback. Controls how the response
* body gets streamed from a non-blocking HTTP connection on the client side.
* @param headers the optional request headers
* @return the response returned by elasticsearch
* @throws IOException in case of a problem or the connection was aborted
* @throws ClientProtocolException in case of an http protocol error
* @throws ResponseException in case elasticsearch responded with a status code that indicated an error
*/
public Response performRequest(String method, String endpoint, Map<String, String> params,
HttpEntity entity, HttpAsyncResponseConsumer<HttpResponse> responseConsumer,
Header... headers) throws Exception {
SyncResponseListener listener = new SyncResponseListener();
performRequest(method, endpoint, params, entity, consumer, listener, headers);
performRequest(method, endpoint, params, entity, responseConsumer, listener, headers);
return listener.get();
}
@ -225,7 +250,7 @@ public final class RestClient implements Closeable {
*/
public void performRequest(String method, String endpoint, Map<String, String> params,
HttpEntity entity, ResponseListener responseListener, Header... headers) {
HttpAsyncResponseConsumer<HttpResponse> responseConsumer = HttpAsyncMethods.createConsumer();
HttpAsyncResponseConsumer<HttpResponse> responseConsumer = new HeapBufferedAsyncResponseConsumer();
performRequest(method, endpoint, params, entity, responseConsumer, responseListener, headers);
}
@ -241,7 +266,8 @@ public final class RestClient implements Closeable {
* @param endpoint the path of the request (without host and port)
* @param params the query_string parameters
* @param entity the body of the request, null if not applicable
* @param responseConsumer the {@link HttpAsyncResponseConsumer} callback
* @param responseConsumer the {@link HttpAsyncResponseConsumer} callback. Controls how the response
* body gets streamed from a non-blocking HTTP connection on the client side.
* @param responseListener the {@link ResponseListener} to notify when the request is completed or fails
* @param headers the optional request headers
*/

View File

@ -0,0 +1,104 @@
/*
* Licensed to Elasticsearch under one or more contributor
* license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright
* ownership. Elasticsearch licenses this file to you under
* the Apache License, Version 2.0 (the "License"); you may
* not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.elasticsearch.client;
import org.apache.http.ContentTooLongException;
import org.apache.http.HttpResponse;
import org.apache.http.ProtocolVersion;
import org.apache.http.StatusLine;
import org.apache.http.entity.BasicHttpEntity;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.message.BasicHttpResponse;
import org.apache.http.message.BasicStatusLine;
import org.apache.http.nio.ContentDecoder;
import org.apache.http.nio.IOControl;
import org.apache.http.protocol.HttpContext;
import static org.elasticsearch.client.HeapBufferedAsyncResponseConsumer.DEFAULT_BUFFER_LIMIT;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
public class HeapBufferedAsyncResponseConsumerTests extends RestClientTestCase {
//maximum buffer that this test ends up allocating is 50MB
private static final int MAX_TEST_BUFFER_SIZE = 50 * 1024 * 1024;
public void testResponseProcessing() throws Exception {
ContentDecoder contentDecoder = mock(ContentDecoder.class);
IOControl ioControl = mock(IOControl.class);
HttpContext httpContext = mock(HttpContext.class);
HeapBufferedAsyncResponseConsumer consumer = spy(new HeapBufferedAsyncResponseConsumer());
ProtocolVersion protocolVersion = new ProtocolVersion("HTTP", 1, 1);
StatusLine statusLine = new BasicStatusLine(protocolVersion, 200, "OK");
HttpResponse httpResponse = new BasicHttpResponse(statusLine);
httpResponse.setEntity(new StringEntity("test"));
//everything goes well
consumer.responseReceived(httpResponse);
consumer.consumeContent(contentDecoder, ioControl);
consumer.responseCompleted(httpContext);
verify(consumer).releaseResources();
verify(consumer).buildResult(httpContext);
assertTrue(consumer.isDone());
assertSame(httpResponse, consumer.getResult());
consumer.responseCompleted(httpContext);
verify(consumer, times(1)).releaseResources();
verify(consumer, times(1)).buildResult(httpContext);
}
public void testDefaultBufferLimit() throws Exception {
HeapBufferedAsyncResponseConsumer consumer = new HeapBufferedAsyncResponseConsumer();
bufferLimitTest(consumer, DEFAULT_BUFFER_LIMIT);
}
public void testConfiguredBufferLimit() throws Exception {
int bufferLimit = randomIntBetween(1, MAX_TEST_BUFFER_SIZE - 100);
HeapBufferedAsyncResponseConsumer consumer = new HeapBufferedAsyncResponseConsumer(bufferLimit);
bufferLimitTest(consumer, bufferLimit);
}
private static void bufferLimitTest(HeapBufferedAsyncResponseConsumer consumer, int bufferLimit) throws Exception {
ProtocolVersion protocolVersion = new ProtocolVersion("HTTP", 1, 1);
StatusLine statusLine = new BasicStatusLine(protocolVersion, 200, "OK");
consumer.onResponseReceived(new BasicHttpResponse(statusLine));
BasicHttpEntity entity = new BasicHttpEntity();
entity.setContentLength(randomInt(bufferLimit));
consumer.onEntityEnclosed(entity, ContentType.APPLICATION_JSON);
entity.setContentLength(randomIntBetween(bufferLimit + 1, MAX_TEST_BUFFER_SIZE));
try {
consumer.onEntityEnclosed(entity, ContentType.APPLICATION_JSON);
} catch(ContentTooLongException e) {
assertEquals("entity content is too long [" + entity.getContentLength() +
"] for the configured buffer limit [" + bufferLimit + "]", e.getMessage());
}
}
}