com.mongodb.internal.operation.AsyncOperationHelper Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of mongodb-driver-core Show documentation
Show all versions of mongodb-driver-core Show documentation
The Java operations layer for the MongoDB Java Driver. Third parties can ' +
'wrap this layer to provide custom higher-level APIs
/*
* Copyright 2008-present 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.internal.operation;
import com.mongodb.Function;
import com.mongodb.MongoException;
import com.mongodb.MongoNamespace;
import com.mongodb.ReadPreference;
import com.mongodb.ServerAddress;
import com.mongodb.assertions.Assertions;
import com.mongodb.internal.async.AsyncBatchCursor;
import com.mongodb.internal.async.SingleResultCallback;
import com.mongodb.internal.async.function.AsyncCallbackBiFunction;
import com.mongodb.internal.async.function.AsyncCallbackFunction;
import com.mongodb.internal.async.function.AsyncCallbackSupplier;
import com.mongodb.internal.async.function.RetryState;
import com.mongodb.internal.async.function.RetryingAsyncCallbackSupplier;
import com.mongodb.internal.binding.AsyncConnectionSource;
import com.mongodb.internal.binding.AsyncReadBinding;
import com.mongodb.internal.binding.AsyncWriteBinding;
import com.mongodb.internal.binding.ReferenceCounted;
import com.mongodb.internal.connection.AsyncConnection;
import com.mongodb.internal.connection.OperationContext;
import com.mongodb.internal.connection.QueryResult;
import com.mongodb.internal.operation.retry.AttachmentKeys;
import com.mongodb.internal.validator.NoOpFieldNameValidator;
import com.mongodb.lang.Nullable;
import org.bson.BsonDocument;
import org.bson.BsonValue;
import org.bson.FieldNameValidator;
import org.bson.codecs.BsonDocumentCodec;
import org.bson.codecs.Decoder;
import java.util.Collections;
import java.util.List;
import static com.mongodb.assertions.Assertions.assertNotNull;
import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
import static com.mongodb.internal.operation.CommandOperationHelper.CommandCreator;
import static com.mongodb.internal.operation.CommandOperationHelper.addRetryableWriteErrorLabel;
import static com.mongodb.internal.operation.CommandOperationHelper.initialRetryState;
import static com.mongodb.internal.operation.CommandOperationHelper.isRetryWritesEnabled;
import static com.mongodb.internal.operation.CommandOperationHelper.logRetryExecute;
import static com.mongodb.internal.operation.CommandOperationHelper.transformWriteException;
import static com.mongodb.internal.operation.OperationHelper.cursorDocumentToQueryResult;
import static com.mongodb.internal.operation.WriteConcernHelper.throwOnWriteConcernError;
final class AsyncOperationHelper {
interface AsyncCallableWithConnection {
void call(@Nullable AsyncConnection connection, @Nullable Throwable t);
}
interface AsyncCallableWithSource {
void call(@Nullable AsyncConnectionSource source, @Nullable Throwable t);
}
interface CommandWriteTransformerAsync {
/**
* Yield an appropriate result object for the input object.
*
* @param t the input object
* @return the function result
*/
@Nullable
R apply(T t, AsyncConnection connection);
}
interface CommandReadTransformerAsync {
/**
* Yield an appropriate result object for the input object.
*
* @param t the input object
* @return the function result
*/
@Nullable
R apply(T t, AsyncConnectionSource source, AsyncConnection connection);
}
static void withAsyncReadConnectionSource(final AsyncReadBinding binding, final AsyncCallableWithSource callable) {
binding.getReadConnectionSource(errorHandlingCallback(new AsyncCallableWithSourceCallback(callable), OperationHelper.LOGGER));
}
static void withAsyncConnection(final AsyncWriteBinding binding, final AsyncCallableWithConnection callable) {
binding.getWriteConnectionSource(errorHandlingCallback(new AsyncCallableWithConnectionCallback(callable), OperationHelper.LOGGER));
}
/**
* @see #withAsyncSuppliedResource(AsyncCallbackSupplier, boolean, SingleResultCallback, AsyncCallbackFunction)
*/
static void withAsyncSourceAndConnection(final AsyncCallbackSupplier sourceSupplier,
final boolean wrapConnectionSourceException, final SingleResultCallback callback,
final AsyncCallbackBiFunction asyncFunction)
throws OperationHelper.ResourceSupplierInternalException {
SingleResultCallback errorHandlingCallback = errorHandlingCallback(callback, OperationHelper.LOGGER);
withAsyncSuppliedResource(sourceSupplier, wrapConnectionSourceException, errorHandlingCallback,
(source, sourceReleasingCallback) ->
withAsyncSuppliedResource(source::getConnection, wrapConnectionSourceException, sourceReleasingCallback,
(connection, connectionAndSourceReleasingCallback) ->
asyncFunction.apply(source, connection, connectionAndSourceReleasingCallback)));
}
static void withAsyncSuppliedResource(final AsyncCallbackSupplier resourceSupplier,
final boolean wrapSourceConnectionException, final SingleResultCallback callback,
final AsyncCallbackFunction function) throws OperationHelper.ResourceSupplierInternalException {
SingleResultCallback errorHandlingCallback = errorHandlingCallback(callback, OperationHelper.LOGGER);
resourceSupplier.get((resource, supplierException) -> {
if (supplierException != null) {
if (wrapSourceConnectionException) {
supplierException = new OperationHelper.ResourceSupplierInternalException(supplierException);
}
errorHandlingCallback.onResult(null, supplierException);
} else {
Assertions.assertNotNull(resource);
AsyncCallbackSupplier curriedFunction = c -> function.apply(resource, c);
curriedFunction.whenComplete(resource::release).get(errorHandlingCallback);
}
});
}
static void withAsyncConnectionSourceCallableConnection(final AsyncConnectionSource source,
final AsyncCallableWithConnection callable) {
source.getConnection((connection, t) -> {
source.release();
if (t != null) {
callable.call(null, t);
} else {
callable.call(connection, null);
}
});
}
static void withAsyncConnectionSource(final AsyncConnectionSource source, final AsyncCallableWithSource callable) {
callable.call(source, null);
}
static void executeRetryableReadAsync(
final AsyncReadBinding binding,
final String database,
final CommandCreator commandCreator,
final Decoder decoder,
final CommandReadTransformerAsync transformer,
final boolean retryReads,
final SingleResultCallback callback) {
executeRetryableReadAsync(binding, binding::getReadConnectionSource, database, commandCreator, decoder, transformer, retryReads,
callback);
}
static void executeRetryableReadAsync(
final AsyncReadBinding binding,
final AsyncCallbackSupplier sourceAsyncSupplier,
final String database,
final CommandCreator commandCreator,
final Decoder decoder,
final CommandReadTransformerAsync transformer,
final boolean retryReads,
final SingleResultCallback callback) {
RetryState retryState = initialRetryState(retryReads);
binding.retain();
AsyncCallbackSupplier asyncRead = decorateReadWithRetriesAsync(retryState, binding.getOperationContext(),
(AsyncCallbackSupplier) funcCallback ->
withAsyncSourceAndConnection(sourceAsyncSupplier, false, funcCallback,
(source, connection, releasingCallback) -> {
if (retryState.breakAndCompleteIfRetryAnd(
() -> !OperationHelper.canRetryRead(source.getServerDescription(),
binding.getSessionContext()),
releasingCallback)) {
return;
}
createReadCommandAndExecuteAsync(retryState, binding, source,
database, commandCreator,
decoder, transformer,
connection,
releasingCallback);
})
).whenComplete(binding::release);
asyncRead.get(errorHandlingCallback(callback, OperationHelper.LOGGER));
}
static void executeCommandAsync(final AsyncWriteBinding binding,
final String database,
final BsonDocument command,
final AsyncConnection connection,
final CommandWriteTransformerAsync transformer,
final SingleResultCallback callback) {
Assertions.notNull("binding", binding);
SingleResultCallback addingRetryableLabelCallback = addingRetryableLabelCallback(callback,
connection.getDescription().getMaxWireVersion());
connection.commandAsync(database, command, new NoOpFieldNameValidator(), ReadPreference.primary(), new BsonDocumentCodec(),
binding, transformingWriteCallback(transformer, connection, addingRetryableLabelCallback));
}
static void executeRetryableWriteAsync(
final AsyncWriteBinding binding,
final String database,
@Nullable final ReadPreference readPreference,
final FieldNameValidator fieldNameValidator,
final Decoder commandResultDecoder,
final CommandCreator commandCreator,
final CommandWriteTransformerAsync transformer,
final Function retryCommandModifier,
final SingleResultCallback callback) {
RetryState retryState = initialRetryState(true);
binding.retain();
AsyncCallbackSupplier asyncWrite = decorateWriteWithRetriesAsync(retryState, binding.getOperationContext(),
(AsyncCallbackSupplier) funcCallback -> {
boolean firstAttempt = retryState.isFirstAttempt();
if (!firstAttempt && binding.getSessionContext().hasActiveTransaction()) {
binding.getSessionContext().clearTransactionContext();
}
withAsyncSourceAndConnection(binding::getWriteConnectionSource, true, funcCallback,
(source, connection, releasingCallback) -> {
int maxWireVersion = connection.getDescription().getMaxWireVersion();
SingleResultCallback addingRetryableLabelCallback = firstAttempt
? releasingCallback
: addingRetryableLabelCallback(releasingCallback, maxWireVersion);
if (retryState.breakAndCompleteIfRetryAnd(() -> !OperationHelper.canRetryWrite(connection.getDescription(), binding.getSessionContext()),
addingRetryableLabelCallback)) {
return;
}
BsonDocument command;
try {
command = retryState.attachment(AttachmentKeys.command())
.map(previousAttemptCommand -> {
Assertions.assertFalse(firstAttempt);
return retryCommandModifier.apply(previousAttemptCommand);
}).orElseGet(() -> commandCreator.create(source.getServerDescription(), connection.getDescription()));
// attach `maxWireVersion`, `retryableCommandFlag` ASAP because they are used to check whether we should retry
retryState.attach(AttachmentKeys.maxWireVersion(), maxWireVersion, true)
.attach(AttachmentKeys.retryableCommandFlag(), isRetryWritesEnabled(command), true)
.attach(AttachmentKeys.commandDescriptionSupplier(), command::getFirstKey, false)
.attach(AttachmentKeys.command(), command, false);
} catch (Throwable t) {
addingRetryableLabelCallback.onResult(null, t);
return;
}
connection.commandAsync(database, command, fieldNameValidator, readPreference, commandResultDecoder, binding,
transformingWriteCallback(transformer, connection, addingRetryableLabelCallback));
});
}).whenComplete(binding::release);
asyncWrite.get(exceptionTransformingCallback(errorHandlingCallback(callback, OperationHelper.LOGGER)));
}
static void createReadCommandAndExecuteAsync(
final RetryState retryState,
final AsyncReadBinding binding,
final AsyncConnectionSource source,
final String database,
final CommandCreator commandCreator,
final Decoder decoder,
final CommandReadTransformerAsync transformer,
final AsyncConnection connection,
final SingleResultCallback callback) {
BsonDocument command;
try {
command = commandCreator.create(source.getServerDescription(), connection.getDescription());
retryState.attach(AttachmentKeys.commandDescriptionSupplier(), command::getFirstKey, false);
} catch (IllegalArgumentException e) {
callback.onResult(null, e);
return;
}
connection.commandAsync(database, command, new NoOpFieldNameValidator(), source.getReadPreference(), decoder,
binding, transformingReadCallback(transformer, source, connection, callback));
}
static AsyncCallbackSupplier decorateReadWithRetriesAsync(final RetryState retryState, final OperationContext operationContext,
final AsyncCallbackSupplier asyncReadFunction) {
return new RetryingAsyncCallbackSupplier<>(retryState, CommandOperationHelper::chooseRetryableReadException,
CommandOperationHelper::shouldAttemptToRetryRead, callback -> {
logRetryExecute(retryState, operationContext);
asyncReadFunction.get(callback);
});
}
static AsyncCallbackSupplier decorateWriteWithRetriesAsync(final RetryState retryState, final OperationContext operationContext,
final AsyncCallbackSupplier asyncWriteFunction) {
return new RetryingAsyncCallbackSupplier<>(retryState, CommandOperationHelper::chooseRetryableWriteException,
CommandOperationHelper::shouldAttemptToRetryWrite, callback -> {
logRetryExecute(retryState, operationContext);
asyncWriteFunction.get(callback);
});
}
static CommandWriteTransformerAsync writeConcernErrorTransformerAsync() {
return (result, connection) -> {
assertNotNull(result);
throwOnWriteConcernError(result, connection.getDescription().getServerAddress(), connection.getDescription().getMaxWireVersion());
return null;
};
}
static AsyncBatchCursor createEmptyAsyncBatchCursor(final MongoNamespace namespace, final ServerAddress serverAddress) {
return new AsyncSingleBatchQueryCursor<>(new QueryResult<>(namespace, Collections.emptyList(), 0L, serverAddress));
}
static AsyncBatchCursor cursorDocumentToAsyncBatchCursor(final BsonDocument cursorDocument, final Decoder decoder,
final BsonValue comment, final AsyncConnectionSource source, final AsyncConnection connection, final int batchSize) {
return new AsyncQueryBatchCursor<>(cursorDocumentToQueryResult(cursorDocument,
source.getServerDescription().getAddress()),
0, batchSize, 0, decoder, comment, source, connection, cursorDocument);
}
static SingleResultCallback releasingCallback(final SingleResultCallback wrapped, final AsyncConnection connection) {
return new ReferenceCountedReleasingWrappedCallback<>(wrapped, Collections.singletonList(connection));
}
static SingleResultCallback exceptionTransformingCallback(final SingleResultCallback callback) {
return (result, t) -> {
if (t != null) {
if (t instanceof MongoException) {
callback.onResult(null, transformWriteException((MongoException) t));
} else {
callback.onResult(null, t);
}
} else {
callback.onResult(result, null);
}
};
}
private static SingleResultCallback transformingWriteCallback(final CommandWriteTransformerAsync transformer,
final AsyncConnection connection, final SingleResultCallback callback) {
return (result, t) -> {
if (t != null) {
callback.onResult(null, t);
} else {
R transformedResult;
try {
transformedResult = transformer.apply(assertNotNull(result), connection);
} catch (Throwable e) {
callback.onResult(null, e);
return;
}
callback.onResult(transformedResult, null);
}
};
}
private static class AsyncCallableWithConnectionCallback implements SingleResultCallback {
private final AsyncCallableWithConnection callable;
AsyncCallableWithConnectionCallback(final AsyncCallableWithConnection callable) {
this.callable = callable;
}
@Override
public void onResult(@Nullable final AsyncConnectionSource source, @Nullable final Throwable t) {
if (t != null) {
callable.call(null, t);
} else {
withAsyncConnectionSourceCallableConnection(Assertions.assertNotNull(source), callable);
}
}
}
private static class AsyncCallableWithSourceCallback implements SingleResultCallback {
private final AsyncCallableWithSource callable;
AsyncCallableWithSourceCallback(final AsyncCallableWithSource callable) {
this.callable = callable;
}
@Override
public void onResult(@Nullable final AsyncConnectionSource source, @Nullable final Throwable t) {
if (t != null) {
callable.call(null, t);
} else {
withAsyncConnectionSource(Assertions.assertNotNull(source), callable);
}
}
}
private static class ReferenceCountedReleasingWrappedCallback implements SingleResultCallback {
private final SingleResultCallback wrapped;
private final List referenceCounted;
ReferenceCountedReleasingWrappedCallback(final SingleResultCallback wrapped,
final List referenceCounted) {
this.wrapped = wrapped;
this.referenceCounted = Assertions.notNull("referenceCounted", referenceCounted);
}
@Override
public void onResult(@Nullable final T result, @Nullable final Throwable t) {
for (ReferenceCounted cur : referenceCounted) {
if (cur != null) {
cur.release();
}
}
wrapped.onResult(result, t);
}
}
private static SingleResultCallback addingRetryableLabelCallback(final SingleResultCallback callback,
final int maxWireVersion) {
return (result, t) -> {
if (t != null) {
if (t instanceof MongoException) {
addRetryableWriteErrorLabel((MongoException) t, maxWireVersion);
}
callback.onResult(null, t);
} else {
callback.onResult(result, null);
}
};
}
private static SingleResultCallback transformingReadCallback(final CommandReadTransformerAsync transformer,
final AsyncConnectionSource source, final AsyncConnection connection, final SingleResultCallback callback) {
return (result, t) -> {
if (t != null) {
callback.onResult(null, t);
} else {
R transformedResult;
try {
transformedResult = transformer.apply(assertNotNull(result), source, connection);
} catch (Throwable e) {
callback.onResult(null, e);
return;
}
callback.onResult(transformedResult, null);
}
};
}
private AsyncOperationHelper() {
}
}