com.mongodb.connection.WriteCommandProtocol Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of mongo-java-driver Show documentation
Show all versions of mongo-java-driver Show documentation
The MongoDB Java Driver uber-artifact, containing mongodb-driver, mongodb-driver-core, and bson
The newest version!
/*
* Copyright (c) 2008-2015 MongoDB, Inc.
*
* Licensed 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 com.mongodb.connection;
import com.mongodb.MongoBulkWriteException;
import com.mongodb.MongoNamespace;
import com.mongodb.WriteConcern;
import com.mongodb.async.SingleResultCallback;
import com.mongodb.bulk.BulkWriteResult;
import com.mongodb.bulk.WriteRequest;
import com.mongodb.event.CommandListener;
import com.mongodb.internal.connection.IndexMap;
import org.bson.BsonDocument;
import org.bson.codecs.BsonDocumentCodec;
import static com.mongodb.connection.ByteBufBsonDocument.createOne;
import static com.mongodb.connection.ProtocolHelper.getCommandFailureException;
import static com.mongodb.connection.ProtocolHelper.getMessageSettings;
import static com.mongodb.connection.ProtocolHelper.sendCommandFailedEvent;
import static com.mongodb.connection.ProtocolHelper.sendCommandStartedEvent;
import static com.mongodb.connection.ProtocolHelper.sendCommandSucceededEvent;
import static com.mongodb.connection.WriteCommandResultHelper.getBulkWriteException;
import static com.mongodb.connection.WriteCommandResultHelper.getBulkWriteResult;
import static java.lang.String.format;
/**
* A base class for implementations of the bulk write commands.
*/
abstract class WriteCommandProtocol implements Protocol {
private final MongoNamespace namespace;
private final boolean ordered;
private final WriteConcern writeConcern;
private CommandListener commandListener;
/**
* Construct an instance.
*
* @param namespace the namespace
* @param ordered whether the inserts are ordered
* @param writeConcern the write concern
*/
public WriteCommandProtocol(final MongoNamespace namespace, final boolean ordered, final WriteConcern writeConcern) {
this.namespace = namespace;
this.ordered = ordered;
this.writeConcern = writeConcern;
}
@Override
public void setCommandListener(final CommandListener commandListener) {
this.commandListener = commandListener;
}
/**
* Gets the write concern.
*
* @return the write concern
*/
public WriteConcern getWriteConcern() {
return writeConcern;
}
@Override
public BulkWriteResult execute(final InternalConnection connection) {
BaseWriteCommandMessage message = createRequestMessage(getMessageSettings(connection.getDescription()));
long startTimeNanos = System.nanoTime();
try {
BulkWriteBatchCombiner bulkWriteBatchCombiner = new BulkWriteBatchCombiner(connection.getDescription().getServerAddress(),
ordered, writeConcern);
int batchNum = 0;
int currentRangeStartIndex = 0;
do {
batchNum++;
startTimeNanos = System.nanoTime();
BaseWriteCommandMessage nextMessage = sendMessage(connection, message, batchNum);
int itemCount = nextMessage != null ? message.getItemCount() - nextMessage.getItemCount() : message.getItemCount();
IndexMap indexMap = IndexMap.create(currentRangeStartIndex, itemCount);
BsonDocument result = receiveMessage(connection, message);
if (nextMessage != null || batchNum > 1) {
if (getLogger().isDebugEnabled()) {
getLogger().debug(format("Received response for batch %d", batchNum));
}
}
if (WriteCommandResultHelper.hasError(result)) {
MongoBulkWriteException bulkWriteException = getBulkWriteException(getType(), result,
connection.getDescription().getServerAddress());
bulkWriteBatchCombiner.addErrorResult(bulkWriteException, indexMap);
} else {
bulkWriteBatchCombiner.addResult(getBulkWriteResult(getType(), result), indexMap);
}
if (commandListener != null) {
sendCommandSucceededEvent(message, message.getCommandName(), result, connection.getDescription(),
startTimeNanos, commandListener);
}
currentRangeStartIndex += itemCount;
message = nextMessage;
} while (message != null && !bulkWriteBatchCombiner.shouldStopSendingMoreBatches());
return bulkWriteBatchCombiner.getResult();
} catch (MongoBulkWriteException e) {
throw e;
} catch (RuntimeException e) {
if (commandListener != null) {
sendCommandFailedEvent(message, message.getCommandName(), connection.getDescription(), startTimeNanos,
e, commandListener);
}
throw e;
}
}
@Override
public void executeAsync(final InternalConnection connection, final SingleResultCallback callback) {
executeBatchesAsync(connection, createRequestMessage(getMessageSettings(connection.getDescription())),
new BulkWriteBatchCombiner(connection.getDescription().getServerAddress(), ordered, writeConcern), 0, 0,
callback);
}
private void executeBatchesAsync(final InternalConnection connection, final BaseWriteCommandMessage message,
final BulkWriteBatchCombiner bulkWriteBatchCombiner, final int batchNum,
final int currentRangeStartIndex, final SingleResultCallback callback) {
try {
if (message != null && !bulkWriteBatchCombiner.shouldStopSendingMoreBatches()) {
final ByteBufferBsonOutput bsonOutput = new ByteBufferBsonOutput(connection);
final BaseWriteCommandMessage nextMessage = message.encode(bsonOutput);
final int itemCount = nextMessage != null ? message.getItemCount() - nextMessage.getItemCount() : message.getItemCount();
final IndexMap indexMap = IndexMap.create(currentRangeStartIndex, itemCount);
final int nextBatchNum = batchNum + 1;
final int nextRangeStartIndex = currentRangeStartIndex + itemCount;
if (nextBatchNum > 1) {
if (getLogger().isDebugEnabled()) {
getLogger().debug(format("Asynchronously sending batch %d", batchNum));
}
}
sendMessageAsync(connection, message.getId(), bsonOutput, new SingleResultCallback() {
@Override
public void onResult(final BsonDocument result, final Throwable t) {
bsonOutput.close();
if (t != null) {
callback.onResult(null, t);
} else {
if (nextBatchNum > 1) {
if (getLogger().isDebugEnabled()) {
getLogger().debug(format("Asynchronously received response for batch %d", batchNum));
}
}
if (WriteCommandResultHelper.hasError(result)) {
bulkWriteBatchCombiner.addErrorResult(getBulkWriteException(getType(), result,
connection.getDescription().getServerAddress()),
indexMap);
} else {
bulkWriteBatchCombiner.addResult(getBulkWriteResult(getType(), result), indexMap);
}
executeBatchesAsync(connection, nextMessage, bulkWriteBatchCombiner, nextBatchNum, nextRangeStartIndex,
callback);
}
}
});
} else {
if (bulkWriteBatchCombiner.hasErrors()) {
callback.onResult(null, bulkWriteBatchCombiner.getError());
} else {
callback.onResult(bulkWriteBatchCombiner.getResult(), null);
}
}
} catch (Throwable t) {
callback.onResult(null, t);
}
}
protected abstract WriteRequest.Type getType();
protected abstract BaseWriteCommandMessage createRequestMessage(final MessageSettings messageSettings);
private BaseWriteCommandMessage sendMessage(final InternalConnection connection, final BaseWriteCommandMessage message,
final int batchNum) {
ByteBufferBsonOutput bsonOutput = new ByteBufferBsonOutput(connection);
try {
RequestMessage.EncodingMetadata encodingMetadata = message.encodeWithMetadata(bsonOutput);
BaseWriteCommandMessage nextMessage = (BaseWriteCommandMessage) encodingMetadata.getNextMessage();
if (commandListener != null) {
sendCommandStartedEvent(message, namespace.getDatabaseName(), message.getCommandName(),
createOne(bsonOutput, encodingMetadata.getFirstDocumentPosition()),
connection.getDescription(), commandListener);
}
if (nextMessage != null || batchNum > 1) {
if (getLogger().isDebugEnabled()) {
getLogger().debug(format("Sending batch %d", batchNum));
}
}
connection.sendMessage(bsonOutput.getByteBuffers(), message.getId());
return nextMessage;
} finally {
bsonOutput.close();
}
}
private BsonDocument receiveMessage(final InternalConnection connection, final RequestMessage message) {
ResponseBuffers responseBuffers = connection.receiveMessage(message.getId());
try {
ReplyMessage replyMessage = new ReplyMessage(responseBuffers, new BsonDocumentCodec(),
message.getId());
BsonDocument result = replyMessage.getDocuments().get(0);
if (!ProtocolHelper.isCommandOk(result)) {
throw getCommandFailureException(result, connection.getDescription().getServerAddress());
}
return result;
} finally {
responseBuffers.close();
}
}
private void sendMessageAsync(final InternalConnection connection, final int messageId, final ByteBufferBsonOutput buffer,
final SingleResultCallback callback) {
SingleResultCallback receiveCallback = new CommandResultCallback(callback,
new BsonDocumentCodec(),
messageId,
connection.getDescription()
.getServerAddress());
connection.sendMessageAsync(buffer.getByteBuffers(), messageId,
new SendMessageCallback(connection, buffer, messageId, callback,
receiveCallback));
}
/**
* Gets the namespace to execute the protocol in.
*
* @return the namespace
*/
public MongoNamespace getNamespace() {
return namespace;
}
/**
* Gets the logger.
*
* @return the logger
*/
protected abstract com.mongodb.diagnostics.logging.Logger getLogger();
/**
* Gets whether the writes must be executed in order.
*
* @return true if the writes must be executed in order
*/
protected boolean isOrdered() {
return ordered;
}
}