MongoClientService.java

package org.elasticsearch.river.mongodb;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.elasticsearch.ElasticsearchException;
import org.elasticsearch.common.component.AbstractLifecycleComponent;
import org.elasticsearch.common.inject.Inject;
import org.elasticsearch.common.inject.Singleton;
import org.elasticsearch.common.settings.Settings;

import com.google.common.base.Strings;
import com.mongodb.MongoClient;
import com.mongodb.MongoClientOptions;
import com.mongodb.MongoCredential;
import com.mongodb.ServerAddress;

@Singleton
public class MongoClientService extends AbstractLifecycleComponent<MongoClientService> {

    private final Map<ClientCacheKey, MongoClient> mongoClients = new HashMap<>();

    private final Object $lock = new Object[0];

    @Inject
    public MongoClientService(Settings settings) {
        super(settings);
    }

    @Override
    protected void doStart() throws ElasticsearchException {
    }

    @Override
    protected void doStop() throws ElasticsearchException {
    }

    @Override
    protected void doClose() throws ElasticsearchException {
        synchronized ($lock) {
            for (MongoClient mongoClient : mongoClients.values()) {
                mongoClient.close();
            }
            mongoClients.clear();
        }
    }

    public MongoClient getMongoClusterClient(MongoDBRiverDefinition definition) {
        return getMongoShardClient(definition, null);
    }

    /**
     * Get or create a {@link MongoClient} for the given {@code servers}.
     *
     * If a client already exists for the given list of servers with the same credentials and
     * options it will be reused, otherwise a new client will be created.
     */
    public MongoClient getMongoShardClient(MongoDBRiverDefinition definition, List<ServerAddress> shardServers) {
        synchronized ($lock) {
            List<ServerAddress> servers = shardServers != null ? shardServers : definition.getMongoServers();

            List<MongoCredential> mongoCredentials = new ArrayList<>();
            if (!Strings.isNullOrEmpty(definition.getMongoLocalUser()) && !Strings.isNullOrEmpty(definition.getMongoLocalPassword())) {
                mongoCredentials.add(MongoCredential.createCredential(
                        definition.getMongoLocalUser(),
                        !Strings.isNullOrEmpty(definition.getMongoLocalAuthDatabase()) ? definition.getMongoLocalAuthDatabase() : MongoDBRiver.MONGODB_LOCAL_DATABASE,
                        definition.getMongoLocalPassword().toCharArray()));
            }
            if (!Strings.isNullOrEmpty(definition.getMongoAdminUser()) && !Strings.isNullOrEmpty(definition.getMongoAdminPassword())) {
                mongoCredentials.add(MongoCredential.createCredential(
                        definition.getMongoAdminUser(),
                        !Strings.isNullOrEmpty(definition.getMongoAdminAuthDatabase()) ? definition.getMongoAdminAuthDatabase() : MongoDBRiver.MONGODB_ADMIN_DATABASE,
                        definition.getMongoAdminPassword().toCharArray()));
            }
            MongoClientOptions mongoClientOptions = definition.getMongoClientOptions();
            ClientCacheKey cacheKey = new ClientCacheKey(servers, mongoCredentials, mongoClientOptions);

            MongoClient mongoClient = mongoClients.get(cacheKey);
            if (mongoClient != null) {
                return mongoClient;
            }

            logger.info("Creating MongoClient for [{}]", servers);
            mongoClient = new MongoClient(servers, mongoCredentials, mongoClientOptions);
            mongoClients.put(cacheKey, mongoClient);
            return mongoClient;
        }
    }

    static class ClientCacheKey {

        private final List<ServerAddress> servers;
        private final List<MongoCredential> mongoCredentials;
        private final MongoClientOptions mongoClientOptions;

        public ClientCacheKey(List<ServerAddress> servers, List<MongoCredential> mongoCredentials, MongoClientOptions mongoClientOptions) {
            this.servers = servers;
            this.mongoCredentials = mongoCredentials;
            this.mongoClientOptions = mongoClientOptions;
        }

        @Override
        public int hashCode() {
            final int prime = 31;
            int result = 1;
            result = prime * result + ((mongoClientOptions == null) ? 0 : mongoClientOptions.hashCode());
            result = prime * result + ((mongoCredentials == null) ? 0 : mongoCredentials.hashCode());
            result = prime * result + ((servers == null) ? 0 : servers.hashCode());
            return result;
        }

        @Override
        public boolean equals(Object obj) {
            if (this == obj)
                return true;
            if (obj == null)
                return false;
            if (getClass() != obj.getClass())
                return false;
            ClientCacheKey other = (ClientCacheKey) obj;
            if (mongoClientOptions == null) {
                if (other.mongoClientOptions != null)
                    return false;
            } else if (!mongoClientOptions.equals(other.mongoClientOptions))
                return false;
            if (mongoCredentials == null) {
                if (other.mongoCredentials != null)
                    return false;
            } else if (!mongoCredentials.equals(other.mongoCredentials))
                return false;
            if (servers == null) {
                if (other.servers != null)
                    return false;
            } else if (!servers.equals(other.servers))
                return false;
            return true;
        }

    }

}