All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.mongodb.operation.CommandOperationHelper Maven / Gradle / Ivy

Go to download

The MongoDB Java Driver uber-artifact, containing mongodb-driver, mongodb-driver-core, and bson

There is a newer version: 3.12.14
Show newest version
/*
 * 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.operation;

import com.mongodb.Function;
import com.mongodb.MongoCommandException;
import com.mongodb.MongoException;
import com.mongodb.MongoNodeIsRecoveringException;
import com.mongodb.MongoNotPrimaryException;
import com.mongodb.MongoSocketException;
import com.mongodb.MongoWriteConcernException;
import com.mongodb.ReadPreference;
import com.mongodb.async.SingleResultCallback;
import com.mongodb.binding.AsyncConnectionSource;
import com.mongodb.binding.AsyncReadBinding;
import com.mongodb.binding.AsyncWriteBinding;
import com.mongodb.binding.ConnectionSource;
import com.mongodb.binding.ReadBinding;
import com.mongodb.binding.WriteBinding;
import com.mongodb.connection.AsyncConnection;
import com.mongodb.connection.Connection;
import com.mongodb.connection.ConnectionDescription;
import com.mongodb.connection.ServerDescription;
import com.mongodb.internal.operation.WriteConcernHelper;
import com.mongodb.internal.validator.NoOpFieldNameValidator;
import com.mongodb.lang.Nullable;
import com.mongodb.session.SessionContext;
import org.bson.BsonDocument;
import org.bson.FieldNameValidator;
import org.bson.codecs.BsonDocumentCodec;
import org.bson.codecs.Decoder;

import java.util.List;

import static com.mongodb.ReadPreference.primary;
import static com.mongodb.assertions.Assertions.notNull;
import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
import static com.mongodb.operation.OperationHelper.AsyncCallableWithConnectionAndSource;
import static com.mongodb.operation.OperationHelper.CallableWithConnectionAndSource;
import static com.mongodb.operation.OperationHelper.CallableWithSource;
import static com.mongodb.operation.OperationHelper.LOGGER;
import static com.mongodb.operation.OperationHelper.canRetryRead;
import static com.mongodb.operation.OperationHelper.canRetryWrite;
import static com.mongodb.operation.OperationHelper.releasingCallback;
import static com.mongodb.operation.OperationHelper.withConnection;
import static com.mongodb.operation.OperationHelper.withConnectionSource;
import static com.mongodb.operation.OperationHelper.withReleasableConnection;
import static java.lang.String.format;
import static java.util.Arrays.asList;

final class CommandOperationHelper {

    interface CommandReadTransformer {

        /**
         * Yield an appropriate result object for the input object.
         *
         * @param t the input object
         * @return the function result
         */
        R apply(T t, ConnectionSource source, Connection connection);
    }

    interface CommandWriteTransformer {

        /**
         * Yield an appropriate result object for the input object.
         *
         * @param t the input object
         * @return the function result
         */
        R apply(T t, Connection connection);
    }

    interface CommandWriteTransformerAsync {

        /**
         * Yield an appropriate result object for the input object.
         *
         * @param t the input object
         * @return the function result
         */
        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
         */
        R apply(T t, AsyncConnectionSource source, AsyncConnection connection);
    }

    static class IdentityReadTransformer implements CommandReadTransformer {
        @Override
        public T apply(final T t, final ConnectionSource source, final Connection connection) {
            return t;
        }
    }

    static class IdentityWriteTransformer implements CommandWriteTransformer {
        @Override
        public T apply(final T t, final Connection connection) {
            return t;
        }
    }

    static class IdentityWriteTransformerAsync implements CommandWriteTransformerAsync {
        @Override
        public T apply(final T t, final AsyncConnection connection) {
            return t;
        }
    }

    static class IdentityTransformerAsync implements CommandReadTransformerAsync {
        @Override
        public T apply(final T t, final AsyncConnectionSource source, final AsyncConnection connection) {
            return t;
        }
    }

    static CommandWriteTransformer writeConcernErrorTransformer() {
        return new CommandWriteTransformer() {
            @Override
            public Void apply(final BsonDocument result, final Connection connection) {
                WriteConcernHelper.throwOnWriteConcernError(result, connection.getDescription().getServerAddress());
                return null;
            }
        };
    }

    static CommandWriteTransformerAsync writeConcernErrorWriteTransformer() {
        return new CommandWriteTransformerAsync() {
            @Override
            public Void apply(final BsonDocument result, final AsyncConnection connection) {
                WriteConcernHelper.throwOnWriteConcernError(result, connection.getDescription().getServerAddress());
                return null;
            }
        };
    }

    static CommandWriteTransformerAsync writeConcernErrorTransformerAsync() {
        return new CommandWriteTransformerAsync() {
            @Override
            public Void apply(final BsonDocument result, final AsyncConnection connection) {
                WriteConcernHelper.throwOnWriteConcernError(result, connection.getDescription().getServerAddress());
                return null;
            }
        };
    }

    static Function noOpRetryCommandModifier() {
        return new Function() {
            @Override
            public BsonDocument apply(final BsonDocument command) {
                return command;
            }
        };
    }

    interface CommandCreator {
        BsonDocument create(ServerDescription serverDescription, ConnectionDescription connectionDescription);
    }

    /* Read Binding Helpers */

    static BsonDocument executeCommand(final ReadBinding binding, final String database, final CommandCreator commandCreator,
                                       final boolean retryReads) {
        return executeCommand(binding, database, commandCreator, new BsonDocumentCodec(), retryReads);
    }

    static  T executeCommand(final ReadBinding binding, final String database, final CommandCreator commandCreator,
                                final CommandReadTransformer transformer, final boolean retryReads) {
        return executeCommand(binding, database, commandCreator, new BsonDocumentCodec(), transformer, retryReads);
    }

    static  T executeCommand(final ReadBinding binding, final String database, final CommandCreator commandCreator,
                                final Decoder decoder, final boolean retryReads) {
        return executeCommand(binding, database, commandCreator, decoder, new IdentityReadTransformer(), retryReads);
    }

    static  T executeCommand(final ReadBinding binding, final String database, final CommandCreator commandCreator,
                                   final Decoder decoder, final CommandReadTransformer transformer, final boolean retryReads) {
        return withConnectionSource(binding, new CallableWithSource() {
            @Override
            public T call(final ConnectionSource source) {
                return executeCommandWithConnection(binding, source, database, commandCreator, decoder,
                        transformer, retryReads, source.getConnection());
            }
        });
    }

    static  T executeCommandWithConnection(final ReadBinding binding, final ConnectionSource source, final String database,
                                                 final CommandCreator commandCreator, final Decoder decoder,
                                                 final CommandReadTransformer transformer, final boolean retryReads,
                                                 final Connection connection) {
        BsonDocument command = null;
        MongoException exception;
        try {
            command = commandCreator.create(source.getServerDescription(), connection.getDescription());
            return executeCommand(database, command, decoder, source, connection, binding.getReadPreference(), transformer,
                    binding.getSessionContext());
        } catch (MongoException e) {
            exception = e;

            if (!shouldAttemptToRetryRead(retryReads, e)) {
                if (retryReads) {
                    logUnableToRetry(command.getFirstKey(), e);
                }
                throw exception;
            }
        } finally {
            connection.release();
        }

        final MongoException originalException = exception;
        return withReleasableConnection(binding, originalException, new CallableWithConnectionAndSource() {
            @Override
            public T call(final ConnectionSource source, final Connection connection) {
                try {
                    if (!canRetryRead(source.getServerDescription(), connection.getDescription(), binding.getSessionContext())) {
                        throw originalException;
                    }
                    BsonDocument retryCommand = commandCreator.create(source.getServerDescription(), connection.getDescription());
                    logRetryExecute(retryCommand.getFirstKey(), originalException);
                    return executeCommand(database, retryCommand, decoder, source, connection, binding.getReadPreference(), transformer,
                            binding.getSessionContext());
                } finally {
                    connection.release();
                }
            }
        });
    }

    /* Write Binding Helpers */

    static BsonDocument executeCommand(final WriteBinding binding, final String database, final BsonDocument command) {
        return executeCommand(binding, database, command, new IdentityWriteTransformer());
    }

    static  T executeCommand(final WriteBinding binding, final String database, final BsonDocument command,
                                final Decoder decoder) {
        return executeCommand(binding, database, command, decoder, new IdentityWriteTransformer());
    }

    static  T executeCommand(final WriteBinding binding, final String database, final BsonDocument command,
                                final CommandWriteTransformer transformer) {
        return executeCommand(binding, database, command, new BsonDocumentCodec(), transformer);
    }

    static  T executeCommand(final WriteBinding binding, final String database, final BsonDocument command,
                                   final Decoder decoder, final CommandWriteTransformer transformer) {
        return executeCommand(binding, database, command, new NoOpFieldNameValidator(), decoder, transformer);
    }

    static  T executeCommand(final WriteBinding binding, final String database, final BsonDocument command,
                                final Connection connection, final CommandWriteTransformer transformer) {
        return executeCommand(binding, database, command, new BsonDocumentCodec(), connection, transformer);
    }

    static  T executeCommand(final WriteBinding binding, final String database, final BsonDocument command,
                                final Decoder decoder, final Connection connection,
                                final CommandWriteTransformer transformer) {
        notNull("binding", binding);
        return executeWriteCommand(database, command, decoder, connection, primary(), transformer, binding.getSessionContext());
    }

    static  T executeCommand(final WriteBinding binding, final String database, final BsonDocument command,
                                final FieldNameValidator fieldNameValidator, final Decoder decoder,
                                final Connection connection, final CommandWriteTransformer transformer) {
        notNull("binding", binding);
        return executeWriteCommand(database, command, fieldNameValidator, decoder, connection, primary(), transformer,
                binding.getSessionContext());
    }

    static  T executeCommand(final WriteBinding binding, final String database, final BsonDocument command,
                                   final FieldNameValidator fieldNameValidator, final Decoder decoder,
                                   final CommandWriteTransformer transformer) {
        return withReleasableConnection(binding, new CallableWithConnectionAndSource() {
            @Override
            public T call(final ConnectionSource source, final Connection connection) {
                try {
                    return transformer.apply(executeCommand(database, command, fieldNameValidator, decoder,
                            source, connection, primary()), connection);
                } finally {
                    connection.release();
                }
            }
        });
    }

    static BsonDocument executeCommand(final WriteBinding binding, final String database, final BsonDocument command,
                                       final Connection connection) {
        notNull("binding", binding);
        return executeWriteCommand(database, command, new BsonDocumentCodec(), connection, primary(),
                binding.getSessionContext());
    }

    /* Private Read Connection Source Helpers */

    private static  T executeCommand(final String database, final BsonDocument command,
                                        final FieldNameValidator fieldNameValidator, final Decoder decoder,
                                        final ConnectionSource source, final Connection connection,
                                        final ReadPreference readPreference) {
        return executeCommand(database, command, fieldNameValidator, decoder, source, connection,
                readPreference, new IdentityReadTransformer(), source.getSessionContext());
    }

    /* Private Connection Helpers */

    private static  T executeCommand(final String database, final BsonDocument command,
                                           final Decoder decoder, final ConnectionSource source, final Connection connection,
                                           final ReadPreference readPreference,
                                           final CommandReadTransformer transformer, final SessionContext sessionContext) {
        return executeCommand(database, command, new NoOpFieldNameValidator(), decoder, source, connection,
                readPreference, transformer, sessionContext);
    }

    private static  T executeCommand(final String database, final BsonDocument command,
                                           final FieldNameValidator fieldNameValidator, final Decoder decoder,
                                           final ConnectionSource source, final Connection connection, final ReadPreference readPreference,
                                           final CommandReadTransformer transformer, final SessionContext sessionContext) {

        return transformer.apply(connection.command(database, command, fieldNameValidator, readPreference, decoder, sessionContext),
                source, connection);
    }

    /* Private Connection Helpers */

    private static  T executeWriteCommand(final String database, final BsonDocument command,
                                             final Decoder decoder, final Connection connection,
                                             final ReadPreference readPreference, final SessionContext sessionContext) {
        return executeWriteCommand(database, command, new NoOpFieldNameValidator(), decoder, connection,
                readPreference, new IdentityWriteTransformer(), sessionContext);
    }

    private static  T executeWriteCommand(final String database, final BsonDocument command,
                                                final Decoder decoder, final Connection connection,
                                                final ReadPreference readPreference,
                                                final CommandWriteTransformer transformer, final SessionContext sessionContext) {
        return executeWriteCommand(database, command, new NoOpFieldNameValidator(), decoder, connection,
                readPreference, transformer, sessionContext);
    }

    private static  T executeWriteCommand(final String database, final BsonDocument command,
                                                final FieldNameValidator fieldNameValidator, final Decoder decoder,
                                                final Connection connection, final ReadPreference readPreference,
                                                final CommandWriteTransformer transformer, final SessionContext sessionContext) {

        return transformer.apply(connection.command(database, command, fieldNameValidator, readPreference, decoder, sessionContext),
                connection);
    }

    /* Async Read Binding Helpers */

    static void executeCommandAsync(final AsyncReadBinding binding,
                                    final String database,
                                    final CommandCreator commandCreator,
                                    final boolean retryReads,
                                    final SingleResultCallback callback) {
        executeCommandAsync(binding, database, commandCreator, new BsonDocumentCodec(), retryReads, callback);
    }

    static  void executeCommandAsync(final AsyncReadBinding binding,
                                        final String database,
                                        final CommandCreator commandCreator,
                                        final Decoder decoder,
                                        final boolean retryReads,
                                        final SingleResultCallback callback) {
        executeCommandAsync(binding, database, commandCreator, decoder, new IdentityTransformerAsync(), retryReads, callback);
    }

    static  void executeCommandAsync(final AsyncReadBinding binding,
                                        final String database,
                                        final CommandCreator commandCreator,
                                        final CommandReadTransformerAsync transformer,
                                        final boolean retryReads,
                                        final SingleResultCallback callback) {
        executeCommandAsync(binding, database, commandCreator, new BsonDocumentCodec(), transformer, retryReads, callback);
    }

    static  void executeCommandAsync(final AsyncReadBinding binding,
                                           final String database,
                                           final CommandCreator commandCreator,
                                           final Decoder decoder,
                                           final CommandReadTransformerAsync transformer,
                                           final boolean retryReads,
                                           final SingleResultCallback originalCallback) {
        final SingleResultCallback errorHandlingCallback = errorHandlingCallback(originalCallback, LOGGER);
        withConnection(binding, new AsyncCallableWithConnectionAndSource() {
            @Override
            public void call(final AsyncConnectionSource source, final AsyncConnection connection, final Throwable t) {
                if (t != null) {
                    releasingCallback(errorHandlingCallback, source, connection).onResult(null, t);
                } else {
                    executeCommandAsyncWithConnection(binding, source, database, commandCreator, decoder, transformer,
                            retryReads, connection, errorHandlingCallback);
                }
            }
        });
    }

    static  void executeCommandAsync(final AsyncReadBinding binding,
                                           final String database,
                                           final CommandCreator commandCreator,
                                           final Decoder decoder,
                                           final CommandReadTransformerAsync transformer,
                                           final boolean retryReads,
                                           final AsyncConnection connection,
                                           final SingleResultCallback originalCallback) {
        final SingleResultCallback errorHandlingCallback = errorHandlingCallback(originalCallback, LOGGER);
        binding.getReadConnectionSource(new SingleResultCallback() {
            @Override
            public void onResult(final AsyncConnectionSource source, final Throwable t) {
                executeCommandAsyncWithConnection(binding, source, database, commandCreator, decoder, transformer, retryReads,
                        connection, errorHandlingCallback);
            }
        });
    }

    static  void executeCommandAsyncWithConnection(final AsyncReadBinding binding,
                                                         final AsyncConnectionSource source,
                                                         final String database,
                                                         final CommandCreator commandCreator,
                                                         final Decoder decoder,
                                                         final CommandReadTransformerAsync transformer,
                                                         final boolean retryReads,
                                                         final AsyncConnection connection,
                                                         final SingleResultCallback callback) {
        try {
            BsonDocument command = commandCreator.create(source.getServerDescription(), connection.getDescription());
            connection.commandAsync(database, command, new NoOpFieldNameValidator(), binding.getReadPreference(), decoder,
                    binding.getSessionContext(),
                    createCommandCallback(binding, source, connection, database, binding.getReadPreference(),
                            command, commandCreator, new NoOpFieldNameValidator(), decoder, transformer, retryReads, callback));
        } catch (IllegalArgumentException e) {
            connection.release();
            callback.onResult(null, e);
        }
    }

    private static  SingleResultCallback createCommandCallback(final AsyncReadBinding binding,
                                                                        final AsyncConnectionSource oldSource,
                                                                        final AsyncConnection oldConnection,
                                                                        final String database,
                                                                        final ReadPreference readPreference,
                                                                        final BsonDocument originalCommand,
                                                                        final CommandCreator commandCreator,
                                                                        final FieldNameValidator fieldNameValidator,
                                                                        final Decoder commandResultDecoder,
                                                                        final CommandReadTransformerAsync transformer,
                                                                        final boolean retryReads,
                                                                        final SingleResultCallback callback) {
        return new SingleResultCallback() {
            @Override
            public void onResult(final T result, final Throwable originalError) {
                SingleResultCallback releasingCallback = releasingCallback(callback, oldSource, oldConnection);
                if (originalError != null) {
                    checkRetryableException(originalError, releasingCallback);
                } else {
                    try {
                        releasingCallback.onResult(transformer.apply(result, oldSource, oldConnection), null);
                    } catch (Throwable transformError) {
                        checkRetryableException(transformError, releasingCallback);
                    }
                }
            }

            private void checkRetryableException(final Throwable originalError, final SingleResultCallback callback) {
                if (!shouldAttemptToRetryRead(retryReads, originalError)) {
                    if (retryReads) {
                        logUnableToRetry(originalCommand.getFirstKey(), originalError);
                    }
                    callback.onResult(null, originalError);
                } else {
                    oldSource.release();
                    oldConnection.release();
                    retryableCommand(originalError);
                }
            }

            private void retryableCommand(final Throwable originalError) {
                withConnection(binding, new AsyncCallableWithConnectionAndSource() {
                    @Override
                    public void call(final AsyncConnectionSource source, final AsyncConnection connection, final Throwable t) {
                        if (t != null) {
                            callback.onResult(null, originalError);
                        } else if (!canRetryRead(source.getServerDescription(), connection.getDescription(),
                                binding.getSessionContext())) {
                            releasingCallback(callback, source, connection).onResult(null, originalError);
                        } else {
                            BsonDocument retryCommand = commandCreator.create(source.getServerDescription(), connection.getDescription());
                            logRetryExecute(retryCommand.getFirstKey(), originalError);
                            connection.commandAsync(database, retryCommand, fieldNameValidator, readPreference,
                                    commandResultDecoder, binding.getSessionContext(),
                                    new TransformingReadResultCallback(transformer, source, connection,
                                            releasingCallback(callback, source, connection)));
                        }
                    }
                });
            }
        };
    }

    static class TransformingReadResultCallback implements SingleResultCallback {
        private final CommandReadTransformerAsync transformer;
        private final AsyncConnectionSource source;
        private final AsyncConnection connection;
        private final SingleResultCallback callback;

        TransformingReadResultCallback(final CommandReadTransformerAsync transformer, final AsyncConnectionSource source,
                                        final AsyncConnection connection, final SingleResultCallback callback) {
            this.transformer = transformer;
            this.source = source;
            this.connection = connection;
            this.callback = callback;
        }

        @Override
        public void onResult(final T result, final Throwable t) {
            if (t != null) {
                callback.onResult(null, t);
            } else {
                try {
                    R transformedResult = transformer.apply(result, source, connection);
                    callback.onResult(transformedResult, null);
                } catch (Throwable transformError) {
                    callback.onResult(null, transformError);
                }
            }
        }
    }

    /* Async Write Binding Helpers */

    static void executeCommandAsync(final AsyncWriteBinding binding,
                                    final String database,
                                    final BsonDocument command,
                                    final SingleResultCallback callback) {
        executeCommandAsync(binding, database, command, new BsonDocumentCodec(), callback);
    }

    static  void executeCommandAsync(final AsyncWriteBinding binding,
                                        final String database,
                                        final BsonDocument command,
                                        final Decoder decoder,
                                        final SingleResultCallback callback) {
        executeCommandAsync(binding, database, command, decoder, new IdentityWriteTransformerAsync(), callback);
    }

    static  void executeCommandAsync(final AsyncWriteBinding binding,
                                        final String database,
                                        final BsonDocument command,
                                        final CommandWriteTransformerAsync transformer,
                                        final SingleResultCallback callback) {
        executeCommandAsync(binding, database, command, new BsonDocumentCodec(), transformer, callback);
    }

    static  void executeCommandAsync(final AsyncWriteBinding binding,
                                           final String database, final BsonDocument command,
                                           final Decoder decoder,
                                           final CommandWriteTransformerAsync transformer,
                                           final SingleResultCallback callback) {
        executeCommandAsync(binding, database, command, new NoOpFieldNameValidator(), decoder, transformer, callback);
    }

    static  void executeCommandAsync(final AsyncWriteBinding binding,
                                        final String database,
                                        final BsonDocument command,
                                        final Decoder decoder,
                                        final AsyncConnection connection,
                                        final CommandWriteTransformerAsync transformer,
                                        final SingleResultCallback callback) {
        notNull("binding", binding);
        executeCommandAsync(database, command, decoder, connection, primary(), transformer, binding.getSessionContext(),
                callback);
    }

    static  void executeCommandAsync(final AsyncWriteBinding binding,
                                        final String database,
                                        final BsonDocument command,
                                        final FieldNameValidator fieldNameValidator,
                                        final Decoder decoder,
                                        final AsyncConnection connection,
                                        final CommandWriteTransformerAsync transformer,
                                        final SingleResultCallback callback) {
        notNull("binding", binding);
        executeCommandAsync(database, command, fieldNameValidator, decoder, connection, primary(), transformer,
                binding.getSessionContext(), callback);
    }

    static  void executeCommandAsync(final AsyncWriteBinding binding,
                                           final String database, final BsonDocument command,
                                           final FieldNameValidator fieldNameValidator,
                                           final Decoder decoder,
                                           final CommandWriteTransformerAsync transformer,
                                           final SingleResultCallback callback) {
        binding.getWriteConnectionSource(new CommandProtocolExecutingCallback(database, command, fieldNameValidator, decoder,
                primary(), transformer, binding.getSessionContext(), errorHandlingCallback(callback, LOGGER)));
    }

    static void executeCommandAsync(final AsyncWriteBinding binding,
                                    final String database,
                                    final BsonDocument command,
                                    final AsyncConnection connection,
                                    final SingleResultCallback callback) {
        executeCommandAsync(binding, database, command, connection, new IdentityWriteTransformerAsync(), callback);
    }

    static  void executeCommandAsync(final AsyncWriteBinding binding,
                                        final String database,
                                        final BsonDocument command,
                                        final AsyncConnection connection,
                                        final CommandWriteTransformerAsync transformer,
                                        final SingleResultCallback callback) {
        notNull("binding", binding);
        executeCommandAsync(database, command, new BsonDocumentCodec(), connection, primary(), transformer,
                binding.getSessionContext(), callback);
    }

    /* Async Connection Helpers */
    private static  void executeCommandAsync(final String database, final BsonDocument command,
                                                   final Decoder decoder, final AsyncConnection connection,
                                                   final ReadPreference readPreference,
                                                   final CommandWriteTransformerAsync transformer,
                                                   final SessionContext sessionContext,
                                                   final SingleResultCallback callback) {
        connection.commandAsync(database, command, new NoOpFieldNameValidator(), readPreference, decoder, sessionContext,
                new SingleResultCallback() {
                    @Override
                    public void onResult(final D result, final Throwable t) {
                        if (t != null) {
                            callback.onResult(null, t);
                        } else {
                            try {
                                T transformedResult = transformer.apply(result, connection);
                                callback.onResult(transformedResult, null);
                            } catch (Exception e) {
                                callback.onResult(null, e);
                            }
                        }
                    }
                });

    }

    private static  void executeCommandAsync(final String database, final BsonDocument command,
                                                   final FieldNameValidator fieldNameValidator,
                                                   final Decoder decoder, final AsyncConnection connection,
                                                   final ReadPreference readPreference,
                                                   final CommandWriteTransformerAsync transformer,
                                                   final SessionContext sessionContext,
                                                   final SingleResultCallback callback) {
        connection.commandAsync(database, command, fieldNameValidator, readPreference, decoder, sessionContext, true, null, null,
                new SingleResultCallback() {
                    @Override
                    public void onResult(final D result, final Throwable t) {
                        if (t != null) {
                            callback.onResult(null, t);
                        } else {
                            try {
                                T transformedResult = transformer.apply(result, connection);
                                callback.onResult(transformedResult, null);
                            } catch (Exception e) {
                                callback.onResult(null, e);
                            }
                        }
                    }
                });
    }

    /* Retryable write helpers */
    static  R executeRetryableCommand(final WriteBinding binding, final String database, final ReadPreference readPreference,
                                            final FieldNameValidator fieldNameValidator, final Decoder commandResultDecoder,
                                            final CommandCreator commandCreator, final CommandWriteTransformer transformer) {
        return executeRetryableCommand(binding, database, readPreference, fieldNameValidator, commandResultDecoder, commandCreator,
                transformer, noOpRetryCommandModifier());
    }

    static  R executeRetryableCommand(final WriteBinding binding, final String database, final ReadPreference readPreference,
                                            final FieldNameValidator fieldNameValidator, final Decoder commandResultDecoder,
                                            final CommandCreator commandCreator, final CommandWriteTransformer transformer,
                                            final Function retryCommandModifier) {
        return withReleasableConnection(binding, new CallableWithConnectionAndSource() {
            @Override
            public R call(final ConnectionSource source, final Connection connection) {
                BsonDocument command = null;
                MongoException exception;
                try {
                    command = commandCreator.create(source.getServerDescription(), connection.getDescription());
                    return transformer.apply(connection.command(database, command, fieldNameValidator, readPreference,
                            commandResultDecoder, binding.getSessionContext()), connection);
                } catch (MongoException e) {
                    exception = e;
                    if (!shouldAttemptToRetryWrite(command, e)) {
                        if (isRetryWritesEnabled(command)) {
                            logUnableToRetry(command.getFirstKey(), e);
                        }
                        throw exception;
                    }
                } finally {
                    connection.release();
                }

                if (binding.getSessionContext().hasActiveTransaction()) {
                    binding.getSessionContext().unpinServerAddress();
                }
                final BsonDocument originalCommand = command;
                final MongoException originalException = exception;
                return withReleasableConnection(binding, originalException, new CallableWithConnectionAndSource() {
                    @Override
                    public R call(final ConnectionSource source, final Connection connection) {
                        try {
                            if (!canRetryWrite(source.getServerDescription(), connection.getDescription(), binding.getSessionContext())) {
                                throw originalException;
                            }
                            BsonDocument retryCommand = retryCommandModifier.apply(originalCommand);
                            logRetryExecute(retryCommand.getFirstKey(), originalException);
                            return transformer.apply(connection.command(database, retryCommand, fieldNameValidator,
                                    readPreference, commandResultDecoder, binding.getSessionContext()),
                                    connection);
                        } finally {
                            connection.release();
                        }
                    }
                });
            }
        });
    }

    static  void executeRetryableCommand(final AsyncWriteBinding binding, final String database, final ReadPreference readPreference,
                                               final FieldNameValidator fieldNameValidator, final Decoder commandResultDecoder,
                                               final CommandCreator commandCreator,
                                               final CommandWriteTransformerAsync transformer,
                                               final SingleResultCallback originalCallback) {
        executeRetryableCommand(binding, database, readPreference, fieldNameValidator, commandResultDecoder, commandCreator, transformer,
                noOpRetryCommandModifier(), originalCallback);
    }

    static  void executeRetryableCommand(final AsyncWriteBinding binding, final String database, final ReadPreference readPreference,
                                               final FieldNameValidator fieldNameValidator, final Decoder commandResultDecoder,
                                               final CommandCreator commandCreator,
                                               final CommandWriteTransformerAsync transformer,
                                               final Function retryCommandModifier,
                                               final SingleResultCallback originalCallback) {
        final SingleResultCallback errorHandlingCallback = errorHandlingCallback(originalCallback, LOGGER);
        binding.getWriteConnectionSource(new SingleResultCallback() {
            @Override
            public void onResult(final AsyncConnectionSource source, final Throwable t) {
                if (t != null) {
                    errorHandlingCallback.onResult(null, t);
                } else {
                    source.getConnection(new SingleResultCallback() {
                        @Override
                        public void onResult(final AsyncConnection connection, final Throwable t) {
                            if (t != null) {
                                releasingCallback(errorHandlingCallback, source).onResult(null, t);
                            } else {
                                try {
                                    BsonDocument command = commandCreator.create(source.getServerDescription(),
                                            connection.getDescription());
                                    connection.commandAsync(database, command, fieldNameValidator, readPreference,
                                            commandResultDecoder, binding.getSessionContext(),
                                            createCommandCallback(binding, source, connection, database, readPreference,
                                                    command, fieldNameValidator, commandResultDecoder, transformer,
                                                    retryCommandModifier, errorHandlingCallback));
                                } catch (Throwable t1) {
                                    releasingCallback(errorHandlingCallback, source, connection).onResult(null, t1);
                                }
                            }
                        }
                    });
                }
            }
        });
    }

    private static  SingleResultCallback createCommandCallback(final AsyncWriteBinding binding,
                                                                        final AsyncConnectionSource oldSource,
                                                                        final AsyncConnection oldConnection,
                                                                        final String database,
                                                                        final ReadPreference readPreference,
                                                                        final BsonDocument command,
                                                                        final FieldNameValidator fieldNameValidator,
                                                                        final Decoder commandResultDecoder,
                                                                        final CommandWriteTransformerAsync transformer,
                                                                        final Function retryCommandModifier,
                                                                        final SingleResultCallback callback) {
        return new SingleResultCallback() {
            @Override
            public void onResult(final T result, final Throwable originalError) {
                SingleResultCallback releasingCallback = releasingCallback(callback, oldSource, oldConnection);
                if (originalError != null) {
                    checkRetryableException(originalError, releasingCallback);
                } else {
                    try {
                        releasingCallback.onResult(transformer.apply(result, oldConnection), null);
                    } catch (Throwable transformError) {
                        checkRetryableException(transformError, releasingCallback);
                    }
                }
            }

            private void checkRetryableException(final Throwable originalError, final SingleResultCallback releasingCallback) {
                if (!shouldAttemptToRetryWrite(command, originalError)) {
                    if (isRetryWritesEnabled(command)) {
                        logUnableToRetry(command.getFirstKey(), originalError);
                    }
                    releasingCallback.onResult(null, originalError);
                } else {
                    oldConnection.release();
                    oldSource.release();
                    if (binding.getSessionContext().hasActiveTransaction()) {
                        binding.getSessionContext().unpinServerAddress();
                    }
                    retryableCommand(originalError);
                }
            }

            private void retryableCommand(final Throwable originalError) {
                final BsonDocument retryCommand = retryCommandModifier.apply(command);
                logRetryExecute(retryCommand.getFirstKey(), originalError);
                withConnection(binding, new AsyncCallableWithConnectionAndSource() {
                    @Override
                    public void call(final AsyncConnectionSource source, final AsyncConnection connection, final Throwable t) {
                        if (t != null) {
                            callback.onResult(null, originalError);
                        } else if (!canRetryWrite(source.getServerDescription(), connection.getDescription(),
                                binding.getSessionContext())) {
                            releasingCallback(callback, source, connection).onResult(null, originalError);
                        } else {
                            connection.commandAsync(database, retryCommand, fieldNameValidator, readPreference,
                                    commandResultDecoder, binding.getSessionContext(),
                                    new TransformingWriteResultCallback(transformer, connection,
                                            releasingCallback(callback, source, connection)));
                        }
                    }
                });
            }
        };
    }

    static class TransformingWriteResultCallback implements SingleResultCallback {
        private final CommandWriteTransformerAsync transformer;
        private final AsyncConnection connection;
        private final SingleResultCallback callback;

        TransformingWriteResultCallback(final CommandWriteTransformerAsync transformer,
                                   final AsyncConnection connection, final SingleResultCallback callback) {
            this.transformer = transformer;
            this.connection = connection;
            this.callback = callback;
        }

        @Override
        public void onResult(final T result, final Throwable t) {
            if (t != null) {
                callback.onResult(null, t);
            } else {
                try {
                    R transformedResult = transformer.apply(result, connection);
                    callback.onResult(transformedResult, null);
                } catch (Throwable transformError) {
                    callback.onResult(null, transformError);
                }
            }
        }
    }

    private static final List RETRYABLE_ERROR_CODES = asList(6, 7, 89, 91, 189, 9001, 13436, 13435, 11602, 11600, 10107);
    static boolean isRetryableException(final Throwable t) {
        if (!(t instanceof MongoException)) {
            return false;
        }

        if (t instanceof MongoSocketException || t instanceof MongoNotPrimaryException || t instanceof MongoNodeIsRecoveringException) {
            return true;
        }
        String errorMessage = t.getMessage();
        if (t instanceof MongoWriteConcernException) {
            errorMessage = ((MongoWriteConcernException) t).getWriteConcernError().getMessage();
        }
        if (errorMessage.contains("not master") || errorMessage.contains("node is recovering")) {
            return true;
        }
        return RETRYABLE_ERROR_CODES.contains(((MongoException) t).getCode());
    }

    /* Misc operation helpers */

    static void rethrowIfNotNamespaceError(final MongoCommandException e) {
        rethrowIfNotNamespaceError(e, null);
    }

    static  T rethrowIfNotNamespaceError(final MongoCommandException e, final T defaultValue) {
        if (!isNamespaceError(e)) {
            throw e;
        }
        return defaultValue;
    }

    static boolean isNamespaceError(final Throwable t) {
        if (t instanceof MongoCommandException) {
            MongoCommandException e = (MongoCommandException) t;
            return (e.getErrorMessage().contains("ns not found") || e.getErrorCode() == 26);
        } else {
            return false;
        }
    }

    private static class CommandProtocolExecutingCallback implements SingleResultCallback {
        private final String database;
        private final BsonDocument command;
        private final Decoder decoder;
        private final ReadPreference readPreference;
        private final FieldNameValidator fieldNameValidator;
        private final CommandWriteTransformerAsync transformer;
        private final SingleResultCallback callback;
        private final SessionContext sessionContext;

        CommandProtocolExecutingCallback(final String database, final BsonDocument command, final FieldNameValidator fieldNameValidator,
                                         final Decoder decoder, final ReadPreference readPreference,
                                         final CommandWriteTransformerAsync transformer, final SessionContext sessionContext,
                                         final SingleResultCallback callback) {
            this.database = database;
            this.command = command;
            this.fieldNameValidator = fieldNameValidator;
            this.decoder = decoder;
            this.readPreference = readPreference;
            this.transformer = transformer;
            this.sessionContext = sessionContext;
            this.callback = callback;
        }

        @Override
        public void onResult(final AsyncConnectionSource source, final Throwable t) {
            if (t != null) {
                callback.onResult(null, t);
            } else {
                source.getConnection(new SingleResultCallback() {
                    @Override
                    public void onResult(final AsyncConnection connection, final Throwable t) {
                        if (t != null) {
                            callback.onResult(null, t);
                        } else {
                            final SingleResultCallback wrappedCallback = releasingCallback(callback, source, connection);
                            connection.commandAsync(database, command, fieldNameValidator, readPreference, decoder, sessionContext,
                                    new SingleResultCallback() {
                                        @Override
                                        public void onResult(final D response, final Throwable t) {
                                            if (t != null) {
                                                wrappedCallback.onResult(null, t);
                                            } else {
                                                wrappedCallback.onResult(transformer.apply(response, connection), null);
                                            }
                                        }
                                    });
                        }
                    }
                });
            }
        }
    }

    private static boolean shouldAttemptToRetryRead(final boolean retryReadsEnabled, final Throwable exception) {
        return retryReadsEnabled && isRetryableException(exception);
    }

    private static boolean shouldAttemptToRetryWrite(@Nullable final BsonDocument command, final Throwable exception) {
        return isRetryWritesEnabled(command) && isRetryableException(exception);
    }

    private static boolean isRetryWritesEnabled(@Nullable final BsonDocument command) {
        return (command != null && (command.containsKey("txnNumber")
                || command.getFirstKey().equals("commitTransaction") || command.getFirstKey().equals("abortTransaction")));
    }

    static boolean shouldAttemptToRetryWrite(final boolean retryWritesEnabled, final Throwable exception) {
        return retryWritesEnabled && isRetryableException(exception);
    }

    static void logRetryExecute(final String operation, final Throwable originalError) {
        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug(format("Retrying operation %s due to an error \"%s\"", operation, originalError));
        }
    }

    static void logUnableToRetry(final String operation, final Throwable originalError) {
        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug(format("Unable to retry operation %s due to error \"%s\"", operation, originalError));
        }
    }

    private CommandOperationHelper() {
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy