package org.xbib.elasticsearch.action.websocket.bulk;
import org.elasticsearch.ElasticsearchIllegalStateException;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.ActionRequest;
import org.elasticsearch.action.bulk.BulkItemResponse;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.Client;
import org.elasticsearch.common.Nullable;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
import org.elasticsearch.common.util.concurrent.EsExecutors;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentBuilderString;
import org.xbib.elasticsearch.websocket.BaseInteractiveHandler;
import org.xbib.elasticsearch.websocket.InteractiveChannel;
import org.xbib.elasticsearch.websocket.InteractiveRequest;
import java.io.IOException;
import java.util.Queue;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
/**
* The bulk handler is derived from the BulkProcessor, but
* offers explicit flushing and can receive requests from multiple threads
* into a global action queue.
* It supports websocket bulk actions and can write back response to an
* interactive channel.
* The bulk volume is not controlled.
* The default concurrency is 32, the number of actions in a bulk is 100.
*/
public class BulkHandler extends BaseInteractiveHandler {
private final BulkHandler.Listener listener;
private final int concurrentRequests;
private final int bulkActions;
private final TimeValue flushInterval;
private final Semaphore semaphore;
private final ScheduledThreadPoolExecutor scheduler;
private final ScheduledFuture scheduledFuture;
private final AtomicLong executionIdGen = new AtomicLong();
private final static Queue<ActionRequest> bulk = ConcurrentCollections.newQueue();
private volatile boolean closed = false;
@Override
public void handleRequest(final InteractiveRequest request, final InteractiveChannel channel) {
// will be overriden by bulk action
}
/**
* A listener for the execution.
*/
public static interface Listener {
/**
* Callback before the bulk is executed.
*/
void beforeBulk(long executionId, BulkRequest request);
/**
* Callback after a successful execution of bulk request.
*/
void afterBulk(long executionId, BulkRequest request, BulkResponse response);
/**
* Callback after a failed execution of bulk request.
*/
void afterBulk(long executionId, BulkRequest request, Throwable failure);
}
/**
* A listener adapter
*/
class ListenerAdapter implements BulkHandler.Listener {
@Override
public void beforeBulk(long executionId, BulkRequest request) {
}
@Override
public void afterBulk(long executionId, BulkRequest request, BulkResponse response) {
}
@Override
public void afterBulk(long executionId, BulkRequest request, Throwable failure) {
}
}
/**
* A builder used to create a build an instance of a bulk processor.
*/
public static class Builder {
private Settings settings;
private Client client;
private BulkHandler.Listener listener;
private int concurrentRequests = 1;
private int bulkActions = 100;
private TimeValue flushInterval = null;
/**
* Creates a builder of bulk processor with the client to use and the
* listener that will be used to be notified on the completion of bulk
* requests.
*/
public Builder(Client client, BulkHandler.Listener listener) {
this.client = client;
this.listener = listener;
}
/**
* Sets the number of concurrent requests allowed to be executed. A
* value of 0 means that only a single request will be allowed to be
* executed. A value of 1 means 1 concurrent request is allowed to be
* executed while accumulating new bulk requests. Defaults to
* <tt>1</tt>.
*/
public BulkHandler.Builder setConcurrentRequests(int concurrentRequests) {
this.concurrentRequests = concurrentRequests;
return this;
}
/**
* Sets when to flush a new bulk request based on the number of actions
* currently added. Defaults to <tt>1000</tt>. Can be set to <tt>-1</tt>
* to disable it.
*/
public BulkHandler.Builder setBulkActions(int bulkActions) {
this.bulkActions = bulkActions;
return this;
}
/**
* Sets a flush interval flushing *any* bulk actions pending if the
* interval passes. Defaults to not set.
* <p/>
* Note, {@link #setBulkActions(int)} can
* be set to <tt>-1</tt> with the flush interval set allowing for
* complete async processing of bulk actions.
*/
public BulkHandler.Builder setFlushInterval(TimeValue flushInterval) {
this.flushInterval = flushInterval;
return this;
}
/**
* Builds a new bulk processor.
*/
public BulkHandler build() {
return new BulkHandler(settings, client, listener, concurrentRequests, bulkActions, flushInterval);
}
}
public static BulkHandler.Builder builder(Client client, BulkHandler.Listener listener) {
return new BulkHandler.Builder(client, listener);
}
public BulkHandler(Settings settings, Client client) {
super(settings, client);
this.listener = new ListenerAdapter();
this.concurrentRequests = 32;
this.bulkActions = 100;
this.semaphore = new Semaphore(concurrentRequests);
this.flushInterval = null;
this.scheduler = null;
this.scheduledFuture = null;
}
BulkHandler(Settings settings, Client client, BulkHandler.Listener listener, int concurrentRequests, int bulkActions, @Nullable TimeValue flushInterval) {
super(settings, client);
this.listener = listener;
this.concurrentRequests = concurrentRequests;
this.bulkActions = bulkActions;
this.semaphore = new Semaphore(concurrentRequests);
this.flushInterval = flushInterval;
if (flushInterval != null) {
this.scheduler = (ScheduledThreadPoolExecutor) Executors.newScheduledThreadPool(1, EsExecutors.daemonThreadFactory(client.settings(), "websocket_bulk_processor"));
this.scheduler.setExecuteExistingDelayedTasksAfterShutdownPolicy(false);
this.scheduler.setContinueExistingPeriodicTasksAfterShutdownPolicy(false);
this.scheduledFuture = this.scheduler.scheduleWithFixedDelay(new BulkHandler.Flush(), flushInterval.millis(), flushInterval.millis(), TimeUnit.MILLISECONDS);
} else {
this.scheduler = null;
this.scheduledFuture = null;
}
}
public BulkHandler.Listener getListener() {
return listener;
}
/**
* Flushes open bulk actions
*/
public synchronized void flush() {
if (closed) {
return;
}
if (bulk.size() > 0) {
execute(null);
}
}
/**
* Flushes open bulk actions
*/
public synchronized void flush(InteractiveChannel channel) {
if (closed) {
return;
}
if (bulk.size() > 0) {
execute(channel);
}
}
/**
* Closes the processor. If flushing by time is enabled, then its shutdown.
* Any remaining bulk actions are flushed.
*/
public synchronized void close() {
if (closed) {
return;
}
closed = true;
if (this.scheduledFuture != null) {
this.scheduledFuture.cancel(false);
this.scheduler.shutdown();
}
if (bulk.size() > 0) {
execute(null);
}
}
/**
* Adds an {@link IndexRequest} to the list of actions to execute. Follows
* the same behavior of {@link IndexRequest} (for example, if no id is
* provided, one will be generated, or usage of the create flag).
*/
public BulkHandler add(IndexRequest request) {
return add((ActionRequest) request);
}
/**
* Adds an {@link DeleteRequest} to the list of actions to execute.
*/
public BulkHandler add(DeleteRequest request) {
return add((ActionRequest) request);
}
public BulkHandler add(ActionRequest request) {
internalAdd(request);
return this;
}
private synchronized void internalAdd(ActionRequest request) {
bulk.add(request);
executeIfNeeded();
}
private void executeIfNeeded() {
if (closed) {
throw new ElasticsearchIllegalStateException("bulk process already closed");
}
if (!isOverTheLimit()) {
return;
}
execute(null);
}
/**
* Adds an {@link IndexRequest} to the list of actions to execute. Follows
* the same behavior of {@link IndexRequest} (for example, if no id is
* provided, one will be generated, or usage of the create flag).
*/
public BulkHandler add(IndexRequest request, InteractiveChannel channel) {
return add((ActionRequest) request, channel);
}
/**
* Adds an {@link DeleteRequest} to the list of actions to execute.
*/
public BulkHandler add(DeleteRequest request, InteractiveChannel channel) {
return add((ActionRequest) request, channel);
}
public BulkHandler add(ActionRequest request, InteractiveChannel channel) {
internalAdd(request, channel);
return this;
}
private synchronized void internalAdd(ActionRequest request, InteractiveChannel channel) {
bulk.add(request);
executeIfNeeded(channel);
}
private void executeIfNeeded(InteractiveChannel channel) {
if (closed) {
throw new ElasticsearchIllegalStateException("bulk process already closed");
}
if (!isOverTheLimit()) {
return;
}
execute(channel);
}
// (currently) needs to be executed under a lock
private void execute(final InteractiveChannel channel) {
final BulkRequest bulkRequest = new BulkRequest().add(bulk);
bulk.clear();
final long executionId = executionIdGen.incrementAndGet();
if (concurrentRequests == 0) {
// execute in a blocking fashion...
try {
listener.beforeBulk(executionId, bulkRequest);
listener.afterBulk(executionId, bulkRequest, client.bulk(bulkRequest).actionGet());
} catch (Exception e) {
listener.afterBulk(executionId, bulkRequest, e);
}
} else {
try {
semaphore.acquire();
} catch (InterruptedException e) {
listener.afterBulk(executionId, bulkRequest, e);
return;
}
listener.beforeBulk(executionId, bulkRequest);
client.bulk(bulkRequest, new ActionListener<BulkResponse>() {
@Override
public void onResponse(BulkResponse response) {
try {
listener.afterBulk(executionId, bulkRequest, response);
if (channel != null) {
channel.sendResponse("bulkresponse", buildResponse(response));
}
} catch (IOException e) {
logger.error("error while sending bulk response", e);
} finally {
semaphore.release();
}
}
@Override
public void onFailure(Throwable t) {
try {
listener.afterBulk(executionId, bulkRequest, t);
if (channel != null) {
channel.sendResponse("bulkresponse", t);
}
} catch (IOException e) {
logger.error("error while sending bulk response", e);
} finally {
semaphore.release();
}
}
});
}
}
private boolean isOverTheLimit() {
if (bulkActions != -1 && bulk.size() > bulkActions) {
return true;
}
return false;
}
class Flush implements Runnable {
@Override
public void run() {
synchronized (BulkHandler.this) {
if (closed) {
return;
}
if (bulk.size() > 0) {
execute(null);
}
}
}
}
/**
* Taken from the REST bulk action.
*
* @param response the bulk response
* @return a content builder with the response
* @throws IOException
*/
private XContentBuilder buildResponse(BulkResponse response) throws IOException {
XContentBuilder builder = jsonBuilder();
builder.startObject();
builder.field(Fields.TOOK, response.getTookInMillis());
builder.startArray(Fields.ITEMS);
for (BulkItemResponse itemResponse : response) {
builder.startObject();
builder.startObject(itemResponse.getOpType());
builder.field(Fields._INDEX, itemResponse.getIndex());
builder.field(Fields._TYPE, itemResponse.getType());
builder.field(Fields._ID, itemResponse.getId());
long version = itemResponse.getVersion();
if (version != -1) {
builder.field(Fields._VERSION, itemResponse.getVersion());
}
if (itemResponse.isFailed()) {
builder.field(Fields.ERROR, itemResponse.getFailureMessage());
} else {
builder.field(Fields.OK, true);
}
builder.endObject();
builder.endObject();
}
builder.endArray();
builder.endObject();
return builder;
}
static final class Fields {
static final XContentBuilderString ITEMS = new XContentBuilderString("items");
static final XContentBuilderString _INDEX = new XContentBuilderString("_index");
static final XContentBuilderString _TYPE = new XContentBuilderString("_type");
static final XContentBuilderString _ID = new XContentBuilderString("_id");
static final XContentBuilderString ERROR = new XContentBuilderString("error");
static final XContentBuilderString OK = new XContentBuilderString("ok");
static final XContentBuilderString TOOK = new XContentBuilderString("took");
static final XContentBuilderString _VERSION = new XContentBuilderString("_version");
static final XContentBuilderString MATCHES = new XContentBuilderString("matches");
}
}