de.otto.mongodb.profiler.DefaultProfilerService.java Source code

Java tutorial

Introduction

Here is the source code for de.otto.mongodb.profiler.DefaultProfilerService.java

Source

/*
 *  Copyright 2013 Robert Gacki <robert.gacki@cgi.com>
 *
 *  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 de.otto.mongodb.profiler;

import com.mongodb.CommandFailureException;
import com.mongodb.DB;
import com.mongodb.MongoClient;
import com.mongodb.MongoException;
import de.otto.mongodb.profiler.collection.CollectionProfiler;
import de.otto.mongodb.profiler.collection.DefaultCollectionProfiler;
import de.otto.mongodb.profiler.op.DefaultOpProfiler;
import de.otto.mongodb.profiler.op.OpProfiler;
import de.otto.mongodb.profiler.op.analyze.Analyzer;
import de.otto.mongodb.profiler.util.Logger;

import java.util.*;
import java.util.concurrent.locks.ReentrantReadWriteLock;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;

/**
 * The default implementation of a ProfilerService.
 *
 * @author Robert Gacki
 */
public class DefaultProfilerService implements ProfilerService {

    private static final Logger logger = Logger.getLogger(DefaultProfilerService.class);

    /**
     * A factory that creates profiler instances supported by this service implementation.
     *
     * @author Robert Gacki
     */
    public interface ProfilerFactory {

        /**
         * Creates a profiler that profiles database operations.
         *
         * @param database    the database to profile
         * @param threadGroup the ThreadGroup to use if the profiler creates threads
         * @return the new instance
         */
        OpProfiler createOpProfiler(DB database, ThreadGroup threadGroup);

        /**
         * Creates a profiler that profiles collections.
         *
         * @param database    the database to profile
         * @param threadGroup the ThreadGroup to use if the profiler creates threads
         * @return the new instance
         */
        CollectionProfiler createCollectionProfiler(DB database, ThreadGroup threadGroup);

    }

    private final Map<String, DefaultProfiledConnection> connections;
    private final Map<String, Map<String, DefaultProfiledDatabase>> databases;

    private final ReentrantReadWriteLock globalLock;

    private final ThreadGroup profilerThreadGroup;

    private final ProfilerFactory profilerFactory;

    public DefaultProfilerService(final Collection<? extends Analyzer> opProfileAnalyzers) {
        this(new DefaultProfilerFactory(opProfileAnalyzers));
    }

    public DefaultProfilerService(final ProfilerFactory profilerFactory) {
        this.connections = Collections.synchronizedMap(new HashMap<String, DefaultProfiledConnection>());
        this.databases = Collections.synchronizedMap(new HashMap<String, Map<String, DefaultProfiledDatabase>>());
        this.globalLock = new ReentrantReadWriteLock();
        this.profilerThreadGroup = new ThreadGroup("MongoDB-Profiling");
        this.profilerFactory = profilerFactory;
    }

    private void shutdown(final DefaultProfiledConnection connection) {

        if (!globalLock.isWriteLocked()) {
            throw new IllegalStateException("Not write locked!");
        }

        synchronized (connection) {
            final Map<String, DefaultProfiledDatabase> removed = databases.remove(connection.getId());
            if (removed != null) {
                final Iterator<DefaultProfiledDatabase> databases = removed.values().iterator();
                while (databases.hasNext()) {
                    final DefaultProfiledDatabase next = databases.next();
                    next.shutdown();
                    databases.remove();
                }
            }
        }
        connection.shutdown();
    }

    public void shutdown() {
        final ReentrantReadWriteLock.WriteLock lock = globalLock.writeLock();
        lock.lock();
        try {
            final Iterator<Map.Entry<String, DefaultProfiledConnection>> connections = this.connections.entrySet()
                    .iterator();
            while (connections.hasNext()) {
                final Map.Entry<String, DefaultProfiledConnection> next = connections.next();
                connections.remove();
                shutdown(next.getValue());
            }
        } finally {
            lock.unlock();
        }
    }

    @Override
    public List<DefaultProfiledConnection> getConnections() {
        final ReentrantReadWriteLock.ReadLock lock = globalLock.readLock();
        lock.lock();
        try {
            return new ArrayList<>(connections.values());
        } finally {
            lock.unlock();
        }
    }

    @Override
    public ProfiledConnection getConnection(final String id) {
        final ReentrantReadWriteLock.ReadLock lock = globalLock.readLock();
        lock.lock();
        try {
            return connections.get(checkNotNull(id));
        } finally {
            lock.unlock();
        }
    }

    private static String createConnectionId(final String name) throws InvalidConnectionNameException {

        boolean valid = false;
        String result = "";
        for (int i = 0; i < name.length(); i++) {
            char c = name.charAt(i);
            if ((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9')) {
                result += c;
                valid = true;
            } else if (c == ' ' || c == '-') {
                result += '-';
            } else if (c == '_') {
                result += c;
            }
        }

        if (result.isEmpty() || !valid) {
            throw new InvalidConnectionNameException(
                    "Connection name must contain at least a character of a-z, A-Z, 0-9!", name);
        }

        return result;
    }

    @Override
    public synchronized ProfiledConnection createConnection(final String name,
            final ProfiledConnectionConfiguration config) throws InvalidConnectionNameException {

        checkNotNull(config);

        final ReentrantReadWriteLock.WriteLock lock = globalLock.writeLock();
        lock.lock();
        try {

            final String id = createConnectionId(name);

            if (connections.containsKey(id)) {
                throw new InvalidConnectionNameException("Another connection by that name already exists!", name);
            }

            final MongoClient client = new MongoClient(new ArrayList<>(config.getServerAddresses()),
                    config.getClientOptions());
            final DefaultProfiledConnection connection = new DefaultProfiledConnection(id, name, config, client);

            connections.put(id, connection);

            return connection;
        } finally {
            lock.unlock();
        }
    }

    @Override
    public void destroyConnection(final ProfiledConnection connection) {

        checkNotNull(connection);

        final ReentrantReadWriteLock.WriteLock lock = globalLock.writeLock();
        lock.lock();
        try {
            final DefaultProfiledConnection realConnection = this.connections.remove(connection.getId());
            if (realConnection != null) {
                shutdown(realConnection);
            }
        } finally {
            lock.unlock();
        }
    }

    @Override
    public ProfiledDatabase addDatabase(final ProfiledConnection connection, final String name)
            throws DatabaseDoesNotExistException, DatabaseAlreadyConfiguredException, AuthenticationException,
            ConnectivityException {
        return addDatabase(connection, name, null, null);
    }

    @Override
    public ProfiledDatabase addDatabase(final ProfiledConnection connection, final String name,
            final String username, final char[] password) throws DatabaseDoesNotExistException,
            DatabaseAlreadyConfiguredException, AuthenticationException, ConnectivityException {

        final String connectionId = checkNotNull(connection).getId();

        final ReentrantReadWriteLock.WriteLock lock = globalLock.writeLock();
        lock.lock();
        try {

            final DefaultProfiledConnection realConnection = connections.get(connectionId);
            checkArgument(realConnection != null, "Connection [%s] does not exist!", connectionId);

            try {
                if (!containsDb(realConnection, name)) {
                    throw new DatabaseDoesNotExistException(String.format("Database [%s] does not exist!", name),
                            name);
                }
            } catch (MongoException e) {
                throw new ConnectivityException(e.getMessage(), e);
            }

            synchronized (realConnection) {
                Map<String, DefaultProfiledDatabase> databases = this.databases.get(connectionId);

                if (databases != null && databases.containsKey(name)) {
                    throw new DatabaseAlreadyConfiguredException(
                            String.format("Database [%s] is already configured!", name), realConnection, name);
                }

                if (databases == null) {
                    databases = Collections.synchronizedMap(new HashMap<String, DefaultProfiledDatabase>());
                    this.databases.put(connectionId, databases);
                }

                final DB db = realConnection.getClient().getDB(name);

                try {
                    if (username != null && password != null) {
                        if (!db.authenticate(username, password)) {
                            throw new AuthenticationException("Failed to authenticate against database.");
                        }
                    }
                } catch (CommandFailureException e) {
                    throw new AuthenticationException("Failed to authenticate against database.", e);
                }

                try {
                    db.getCollectionNames();
                } catch (MongoException e) {
                    if (e.getCode() == 16550) {
                        throw new AuthenticationException("Database requires authentication.");
                    }
                    throw new ConnectivityException("Failed to test a database command!", e);
                }

                final CollectionProfiler collectionProfiler = profilerFactory.createCollectionProfiler(db,
                        profilerThreadGroup);
                final OpProfiler opProfiler = profilerFactory.createOpProfiler(db, profilerThreadGroup);

                final DefaultProfiledDatabase database = new DefaultProfiledDatabase(db, realConnection,
                        collectionProfiler, opProfiler);

                database.initialize();
                databases.put(name, database);

                return database;
            }
        } finally {
            lock.unlock();
        }
    }

    @Override
    public List<? extends ProfiledDatabase> getDatabases(final ProfiledConnection connection) {

        checkNotNull(connection);

        final ReentrantReadWriteLock.ReadLock lock = globalLock.readLock();
        lock.lock();
        try {
            final Map<String, DefaultProfiledDatabase> databases = this.databases.get(connection.getId());

            if (databases == null) {
                return Collections.emptyList();
            }

            return new ArrayList<>(databases.values());
        } finally {
            lock.unlock();
        }
    }

    @Override
    public ProfiledDatabase getDatabase(final ProfiledConnection connection, final String name) {

        checkNotNull(connection);
        checkNotNull(name);

        final ReentrantReadWriteLock.ReadLock lock = globalLock.readLock();
        lock.lock();
        try {
            final Map<String, DefaultProfiledDatabase> databases = this.databases.get(connection.getId());

            if (databases == null) {
                return null;
            }

            return databases.get(name);
        } finally {
            lock.unlock();
        }
    }

    @Override
    public void removeDatabase(ProfiledDatabase candiate) {

        checkNotNull(candiate);

        final ReentrantReadWriteLock.WriteLock lock = globalLock.writeLock();
        lock.lock();
        try {

            final DefaultProfiledConnection connection = this.connections.get(candiate.getConnection().getId());

            if (connection == null || candiate.getConnection() != connection) {
                return;
            }

            final Map<String, DefaultProfiledDatabase> databases = this.databases.get(connection.getId());

            if (databases == null || databases.isEmpty()) {
                return;
            }

            final DefaultProfiledDatabase database = databases.remove(candiate.getName());

            if (database != null) {
                database.shutdown();
            }

        } finally {
            lock.unlock();
        }
    }

    @Override
    public CollectionProfiler getCollectionProfiler(final ProfiledDatabase database) {

        if (database instanceof DefaultProfiledDatabase) {
            return ((DefaultProfiledDatabase) database).getCollectionProfiler();
        }

        throw new IllegalArgumentException("Invalid database reference!");
    }

    @Override
    public OpProfiler getOpProfiler(final ProfiledDatabase database) {

        if (database instanceof DefaultProfiledDatabase) {
            return ((DefaultProfiledDatabase) database).getOpProfiler();
        }

        throw new IllegalArgumentException("Invalid database reference!");
    }

    private static boolean containsDb(final DefaultProfiledConnection connection, final String name) {

        final List<String> databaseNames;
        try {
            databaseNames = connection.getClient().getDatabaseNames();
        } catch (CommandFailureException e) {
            logger.info(e, "Unable to retrieve database names!");
            return true;
        }

        for (String databaseName : databaseNames) {
            if (name.equals(databaseName)) {
                return true;
            }
        }
        return false;
    }

    private static class DefaultProfilerFactory implements ProfilerFactory {

        private final Collection<? extends Analyzer> opProfileAnalyzers;

        private DefaultProfilerFactory(Collection<? extends Analyzer> opProfileAnalyzers) {
            this.opProfileAnalyzers = checkNotNull(opProfileAnalyzers);
        }

        @Override
        public DefaultOpProfiler createOpProfiler(final DB database, final ThreadGroup threadGroup) {
            return new DefaultOpProfiler(false, database, opProfileAnalyzers, threadGroup);
        }

        @Override
        public DefaultCollectionProfiler createCollectionProfiler(final DB database,
                final ThreadGroup threadGroup) {
            return new DefaultCollectionProfiler(false, database, threadGroup);
        }
    }

    private static class DefaultProfiledConnection implements ProfiledConnection {

        private final String id;
        private final String name;
        private final ProfiledConnectionConfiguration configuration;
        private final MongoClient client;

        public DefaultProfiledConnection(final String id, final String name,
                final ProfiledConnectionConfiguration configuration, final MongoClient client) {
            this.id = checkNotNull(id);
            this.name = checkNotNull(name);
            this.configuration = checkNotNull(configuration);
            this.client = checkNotNull(client);
        }

        @Override
        public String getId() {
            return id;
        }

        @Override
        public String getName() {
            return name;
        }

        @Override
        public ProfiledConnectionConfiguration getConfiguration() {
            return configuration;
        }

        public MongoClient getClient() {
            return client;
        }

        public void shutdown() {
            client.close();
        }
    }

    private static class DefaultProfiledDatabase implements ProfiledDatabase {

        private static void initialize(Profiler... profilers) {
            for (Profiler profiler : profilers) {
                if (profiler instanceof InitializingProfiler) {
                    ((InitializingProfiler) profiler).initialize();
                }
            }
        }

        private static void shutdown(Profiler... profilers) {
            for (Profiler profiler : profilers) {
                if (profiler instanceof AsyncProfiler) {
                    ((AsyncProfiler) profiler).stopProfiling();
                }
                if (profiler instanceof InitializingProfiler) {
                    ((InitializingProfiler) profiler).shutdown();
                }
            }
        }

        private final DB db;
        private final DefaultProfiledConnection connection;
        private final CollectionProfiler collectionProfiler;
        private final OpProfiler opProfiler;

        public DefaultProfiledDatabase(final DB db, final DefaultProfiledConnection connection,
                final CollectionProfiler collectionProfiler, final OpProfiler opProfiler) {
            this.db = checkNotNull(db);
            this.connection = checkNotNull(connection);
            this.collectionProfiler = checkNotNull(collectionProfiler);
            this.opProfiler = checkNotNull(opProfiler);
        }

        @Override
        public String getName() {
            return db.getName();
        }

        @Override
        public DefaultProfiledConnection getConnection() {
            return connection;
        }

        public CollectionProfiler getCollectionProfiler() {
            return collectionProfiler;
        }

        public OpProfiler getOpProfiler() {
            return opProfiler;
        }

        public void initialize() {
            initialize(collectionProfiler, opProfiler);
        }

        public void shutdown() {
            shutdown(collectionProfiler, opProfiler);
        }
    }
}