/*
* Licensed to Elastic Search and Shay Banon under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. Elastic Search 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.action.bulk;
import org.elasticsearch.ElasticSearchException;
import org.elasticsearch.ExceptionsHelper;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.ActionRequest;
import org.elasticsearch.action.TransportActions;
import org.elasticsearch.action.admin.indices.create.CreateIndexRequest;
import org.elasticsearch.action.admin.indices.create.CreateIndexResponse;
import org.elasticsearch.action.admin.indices.create.TransportCreateIndexAction;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.support.BaseAction;
import org.elasticsearch.cluster.ClusterService;
import org.elasticsearch.cluster.ClusterState;
import org.elasticsearch.cluster.metadata.MappingMetaData;
import org.elasticsearch.cluster.routing.GroupShardsIterator;
import org.elasticsearch.cluster.routing.ShardIterator;
import org.elasticsearch.common.UUID;
import org.elasticsearch.common.collect.Lists;
import org.elasticsearch.common.collect.Maps;
import org.elasticsearch.common.collect.Sets;
import org.elasticsearch.common.inject.Inject;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.index.shard.ShardId;
import org.elasticsearch.indices.IndexAlreadyExistsException;
import org.elasticsearch.threadpool.ThreadPool;
import org.elasticsearch.transport.BaseTransportRequestHandler;
import org.elasticsearch.transport.TransportChannel;
import org.elasticsearch.transport.TransportService;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
/**
* @author kimchy (shay.banon)
*/
public class TransportBulkAction extends BaseAction<BulkRequest, BulkResponse> {
private final boolean autoCreateIndex;
private final boolean allowIdGeneration;
private final ClusterService clusterService;
private final TransportShardBulkAction shardBulkAction;
private final TransportCreateIndexAction createIndexAction;
@Inject public TransportBulkAction(Settings settings, ThreadPool threadPool, TransportService transportService, ClusterService clusterService,
TransportShardBulkAction shardBulkAction, TransportCreateIndexAction createIndexAction) {
super(settings, threadPool);
this.clusterService = clusterService;
this.shardBulkAction = shardBulkAction;
this.createIndexAction = createIndexAction;
this.autoCreateIndex = settings.getAsBoolean("action.auto_create_index", true);
this.allowIdGeneration = componentSettings.getAsBoolean("action.allow_id_generation", true);
transportService.registerHandler(TransportActions.BULK, new TransportHandler());
}
@Override protected void doExecute(final BulkRequest bulkRequest, final ActionListener<BulkResponse> listener) {
final long startTime = System.currentTimeMillis();
Set<String> indices = Sets.newHashSet();
for (ActionRequest request : bulkRequest.requests) {
if (request instanceof IndexRequest) {
IndexRequest indexRequest = (IndexRequest) request;
if (!indices.contains(indexRequest.index())) {
indices.add(indexRequest.index());
}
} else if (request instanceof DeleteRequest) {
DeleteRequest deleteRequest = (DeleteRequest) request;
if (!indices.contains(deleteRequest.index())) {
indices.add(deleteRequest.index());
}
}
}
if (autoCreateIndex) {
final AtomicInteger counter = new AtomicInteger(indices.size());
final AtomicBoolean failed = new AtomicBoolean();
for (String index : indices) {
if (!clusterService.state().metaData().hasConcreteIndex(index)) {
createIndexAction.execute(new CreateIndexRequest(index).cause("auto(bulk api)"), new ActionListener<CreateIndexResponse>() {
@Override public void onResponse(CreateIndexResponse result) {
if (counter.decrementAndGet() == 0) {
executeBulk(bulkRequest, startTime, listener);
}
}
@Override public void onFailure(Throwable e) {
if (ExceptionsHelper.unwrapCause(e) instanceof IndexAlreadyExistsException) {
// we have the index, do it
if (counter.decrementAndGet() == 0) {
executeBulk(bulkRequest, startTime, listener);
}
} else if (failed.compareAndSet(false, true)) {
listener.onFailure(e);
}
}
});
} else {
if (counter.decrementAndGet() == 0) {
executeBulk(bulkRequest, startTime, listener);
}
}
}
} else {
executeBulk(bulkRequest, startTime, listener);
}
}
private void executeBulk(final BulkRequest bulkRequest, final long startTime, final ActionListener<BulkResponse> listener) {
ClusterState clusterState = clusterService.state();
for (ActionRequest request : bulkRequest.requests) {
if (request instanceof IndexRequest) {
IndexRequest indexRequest = (IndexRequest) request;
indexRequest.routing(clusterState.metaData().resolveIndexRouting(indexRequest.routing(), indexRequest.index()));
indexRequest.index(clusterState.metaData().concreteIndex(indexRequest.index()));
if (allowIdGeneration) {
if (indexRequest.id() == null) {
indexRequest.id(UUID.randomBase64UUID());
// since we generate the id, change it to CREATE
indexRequest.opType(IndexRequest.OpType.CREATE);
}
}
} else if (request instanceof DeleteRequest) {
DeleteRequest deleteRequest = (DeleteRequest) request;
deleteRequest.index(clusterState.metaData().concreteIndex(deleteRequest.index()));
}
}
final BulkItemResponse[] responses = new BulkItemResponse[bulkRequest.requests.size()];
// first, go over all the requests and create a ShardId -> Operations mapping
Map<ShardId, List<BulkItemRequest>> requestsByShard = Maps.newHashMap();
for (int i = 0; i < bulkRequest.requests.size(); i++) {
ActionRequest request = bulkRequest.requests.get(i);
if (request instanceof IndexRequest) {
IndexRequest indexRequest = (IndexRequest) request;
// handle routing
MappingMetaData mappingMd = clusterState.metaData().index(indexRequest.index()).mapping(indexRequest.type());
if (mappingMd != null) {
try {
indexRequest.processRouting(mappingMd);
} catch (ElasticSearchException e) {
responses[i] = new BulkItemResponse(i, indexRequest.opType().toString().toLowerCase(),
new BulkItemResponse.Failure(indexRequest.index(), indexRequest.type(), indexRequest.id(), e.getDetailedMessage()));
continue;
}
}
ShardId shardId = clusterService.operationRouting().indexShards(clusterState, indexRequest.index(), indexRequest.type(), indexRequest.id(), indexRequest.routing()).shardId();
List<BulkItemRequest> list = requestsByShard.get(shardId);
if (list == null) {
list = Lists.newArrayList();
requestsByShard.put(shardId, list);
}
list.add(new BulkItemRequest(i, request));
} else if (request instanceof DeleteRequest) {
DeleteRequest deleteRequest = (DeleteRequest) request;
MappingMetaData mappingMd = clusterState.metaData().index(deleteRequest.index()).mapping(deleteRequest.type());
if (mappingMd != null && mappingMd.routing().required() && deleteRequest.routing() == null) {
// if routing is required, and no routing on the delete request, we need to broadcast it....
GroupShardsIterator groupShards = clusterService.operationRouting().broadcastDeleteShards(clusterState, deleteRequest.index());
for (ShardIterator shardIt : groupShards) {
List<BulkItemRequest> list = requestsByShard.get(shardIt.shardId());
if (list == null) {
list = Lists.newArrayList();
requestsByShard.put(shardIt.shardId(), list);
}
list.add(new BulkItemRequest(i, request));
}
} else {
ShardId shardId = clusterService.operationRouting().deleteShards(clusterState, deleteRequest.index(), deleteRequest.type(), deleteRequest.id(), deleteRequest.routing()).shardId();
List<BulkItemRequest> list = requestsByShard.get(shardId);
if (list == null) {
list = Lists.newArrayList();
requestsByShard.put(shardId, list);
}
list.add(new BulkItemRequest(i, request));
}
}
}
if (requestsByShard.isEmpty()) {
listener.onResponse(new BulkResponse(responses, System.currentTimeMillis() - startTime));
return;
}
final AtomicInteger counter = new AtomicInteger(requestsByShard.size());
for (Map.Entry<ShardId, List<BulkItemRequest>> entry : requestsByShard.entrySet()) {
final ShardId shardId = entry.getKey();
final List<BulkItemRequest> requests = entry.getValue();
BulkShardRequest bulkShardRequest = new BulkShardRequest(shardId.index().name(), shardId.id(), bulkRequest.refresh(), requests.toArray(new BulkItemRequest[requests.size()]));
bulkShardRequest.replicationType(bulkRequest.replicationType());
bulkShardRequest.consistencyLevel(bulkRequest.consistencyLevel());
shardBulkAction.execute(bulkShardRequest, new ActionListener<BulkShardResponse>() {
@Override public void onResponse(BulkShardResponse bulkShardResponse) {
synchronized (responses) {
for (BulkItemResponse bulkItemResponse : bulkShardResponse.responses()) {
responses[bulkItemResponse.itemId()] = bulkItemResponse;
}
}
if (counter.decrementAndGet() == 0) {
finishHim();
}
}
@Override public void onFailure(Throwable e) {
// create failures for all relevant requests
String message = ExceptionsHelper.detailedMessage(e);
synchronized (responses) {
for (BulkItemRequest request : requests) {
if (request.request() instanceof IndexRequest) {
IndexRequest indexRequest = (IndexRequest) request.request();
responses[request.id()] = new BulkItemResponse(request.id(), indexRequest.opType().toString().toLowerCase(),
new BulkItemResponse.Failure(indexRequest.index(), indexRequest.type(), indexRequest.id(), message));
} else if (request.request() instanceof DeleteRequest) {
DeleteRequest deleteRequest = (DeleteRequest) request.request();
responses[request.id()] = new BulkItemResponse(request.id(), "delete",
new BulkItemResponse.Failure(deleteRequest.index(), deleteRequest.type(), deleteRequest.id(), message));
}
}
}
if (counter.decrementAndGet() == 0) {
finishHim();
}
}
private void finishHim() {
listener.onResponse(new BulkResponse(responses, System.currentTimeMillis() - startTime));
}
});
}
}
class TransportHandler extends BaseTransportRequestHandler<BulkRequest> {
@Override public BulkRequest newInstance() {
return new BulkRequest();
}
@Override public void messageReceived(final BulkRequest request, final TransportChannel channel) throws Exception {
// no need to use threaded listener, since we just send a response
request.listenerThreaded(false);
execute(request, new ActionListener<BulkResponse>() {
@Override public void onResponse(BulkResponse result) {
try {
channel.sendResponse(result);
} catch (Exception e) {
onFailure(e);
}
}
@Override public void onFailure(Throwable e) {
try {
channel.sendResponse(e);
} catch (Exception e1) {
logger.warn("Failed to send error response for action [" + TransportActions.BULK + "] and request [" + request + "]", e1);
}
}
});
}
@Override public String executor() {
return ThreadPool.Names.SAME;
}
}
}