com.crosstreelabs.cognitio.Main.java Source code

Java tutorial

Introduction

Here is the source code for com.crosstreelabs.cognitio.Main.java

Source

/*
 * Copyright 2015 Crosstree Labs.
 *
 * 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.crosstreelabs.cognitio;

import com.beust.jcommander.JCommander;
import com.crosstreelabs.cognitio.api.extension.Binder;
import com.crosstreelabs.cognitio.api.extension.Extension;
import com.crosstreelabs.cognitio.api.extension.FormatHandler;
import com.crosstreelabs.cognitio.api.extension.TransportHandler;
import com.crosstreelabs.cognitio.api.io.Communicator;
import com.crosstreelabs.cognitio.api.io.ConductorCommunicator;
import com.crosstreelabs.cognitio.api.io.LibrarianCommunicator;
import com.crosstreelabs.cognitio.api.io.WorkerCommunicator;
import com.crosstreelabs.cognitio.api.resource.CatalogueEntry;
import com.crosstreelabs.cognitio.api.resource.Host;
import com.crosstreelabs.cognitio.api.resource.Status;
import com.crosstreelabs.cognitio.api.service.CatalogueService;
import com.crosstreelabs.cognitio.api.service.HostService;
import com.crosstreelabs.cognitio.api.service.IndexService;
import com.crosstreelabs.cognitio.api.service.ReferenceService;
import com.crosstreelabs.cognitio.api.worker.Worker;
import com.crosstreelabs.cognitio.common.CustomCommander;
import com.crosstreelabs.cognitio.common.ExtensionLoader;
import com.crosstreelabs.cognitio.common.PropertyFileDefaultProvider;
import com.crosstreelabs.cognitio.conductor.ConductorRunner;
import com.crosstreelabs.cognitio.gumshoe.WorkerRunner;
import com.crosstreelabs.cognitio.gumshoe.util.FormatHandlerSet;
import com.crosstreelabs.cognitio.gumshoe.util.TransportHandlerSet;
import com.crosstreelabs.cognitio.librarian.LibrarianRunner;
import com.google.common.net.InternetDomainName;
import io.mola.galimatias.GalimatiasParseException;
import java.io.File;
import java.io.IOException;
import java.lang.reflect.Type;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import javax.inject.Provider;
import org.apache.commons.io.FileUtils;
import org.injectr.injectr.Injector;
import org.injectr.injectr.StandardInjector;
import org.reflections.Reflections;
import org.reflections.util.ClasspathHelper;
import org.reflections.util.ConfigurationBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class Main {
    private static final Logger LOGGER = LoggerFactory.getLogger(Main.class);
    private static final String VERSION = "0.1.0-SNAPSHOT";

    private JCommander jc;
    private final MasterConfiguration mc = new MasterConfiguration();
    private Injector injector = new StandardInjector();

    public Main(final String[] args) throws Exception {
        // First run
        jc = new CustomCommander(mc);
        jc.setAcceptUnknownOptions(true);
        jc.parseWithoutValidation(args);

        // Second run
        String propertiesFile = getApplicationPropertiesFile();
        jc = new CustomCommander();
        jc.addObject(mc);
        injector.bind(mc);
        Set<Extension> extensions = loadExtensions();
        if (propertiesFile != null) {
            jc.setDefaultProvider(new PropertyFileDefaultProvider(getApplicationPropertiesFile()));
        }
        jc.setColumnSize(80);
        jc.parse(args);

        if (mc.help) {
            jc.usage();
            return;
        }
        if (mc.version) {
            System.out.println(VERSION);
            return;
        }

        // Build the DI stuff
        injector.bind(buildTransportHandlerSet());
        injector.bind(buildFormatHandlerSet());
        //        injector.bindToProvider(Communicator.class, mc.communicatorProvider);
        injector.bind(WorkerCommunicator.class,
                Class.forName("com.crosstreelabs.cognitio.communicator.amqp.service.AMQPWorkerCommunicator"));
        injector.bind(LibrarianCommunicator.class,
                Class.forName("com.crosstreelabs.cognitio.communicator.amqp.service.AMQPLibrarianCommunicator"));
        injector.bindToProvider(ConductorCommunicator.class, mc.communicatorProvider);
        injector.bindToProvider(CatalogueService.class, mc.conductorCatalogueServiceProvider);
        injector.bindToProvider(HostService.class, mc.conductorHostServiceProvider);
        injector.bindToProvider(ReferenceService.class, mc.conductorReferenceServiceProvider);
        injector.bindToProvider(IndexService.class, mc.librarianIndexServiceProvider);
        injector.bindToProvider(Worker.class, mc.workerProvider);

        // Load the seeds
        loadSeeds();

        // Execute the requested component(s)
        if (mc.components.contains("conductor")) {
            ((ConductorRunner) injector.get(ConductorRunner.class)).run();
        }
        if (mc.components.contains("librarian")) {
            ((LibrarianRunner) injector.get(LibrarianRunner.class)).run();
        }
        if (mc.components.contains("worker")) {
            ((WorkerRunner) injector.get(WorkerRunner.class)).run();
        }
    }

    private String getApplicationPropertiesFile() {
        if (mc.configFile != null) {
            return mc.configFile.getAbsolutePath();
        }
        if (mc.homeDir != null) {
            return new File(mc.homeDir, "application.properties").getAbsolutePath();
        }
        File fallback = new File("application.properties");
        if (fallback.exists()) {
            return fallback.getAbsolutePath();
        }
        return null;
    }

    private Set<Extension> loadExtensions()
            throws MalformedURLException, ClassNotFoundException, InstantiationException, IllegalAccessException {
        // Attempt to load the extensions
        File extDir;
        if (mc.extensionsDir != null) {
            extDir = mc.extensionsDir;
        } else if (mc.homeDir != null) {
            extDir = new File(mc.homeDir, "extensions");
        } else {
            extDir = new File("extensions").getAbsoluteFile();
        }
        LOGGER.info("Extension dir: {}", extDir);

        // Load extensions
        ExtensionLoader loader = new ExtensionLoader(new URL[0], Main.class.getClassLoader());
        loader.add(extDir.getAbsolutePath());

        // Load the extensions and give them the opportunity to add command line
        // options
        ConfigurationBuilder rcb = new ConfigurationBuilder();
        rcb.setClassLoaders(new ClassLoader[] { loader, getClass().getClassLoader() });
        rcb.addUrls(ClasspathHelper.forClassLoader(loader, getClass().getClassLoader()));
        Reflections reflections = new Reflections(rcb);
        Set<Class<? extends Extension>> extensions = reflections.getSubTypesOf(Extension.class);
        Set<Extension> result = new HashSet<>();
        for (Class<? extends Extension> ext : extensions) {
            Extension e = ext.newInstance();
            LOGGER.info("Loading extension: {}", e.getClass().getName());
            Object config = e.getConfigurationObject();
            jc.addObject(config);
            injector.bind(config);
            e.bind(new ExtensionBinder(injector));
            result.add(e);
        }
        return result;
    }

    private TransportHandlerSet buildTransportHandlerSet() {
        Collection<Class<? extends TransportHandler>> handlerClasses;
        if (mc.workerTransportHandlers != null && !mc.workerTransportHandlers.isEmpty()) {
            handlerClasses = mc.workerTransportHandlers;
        } else {
            ConfigurationBuilder rcb = new ConfigurationBuilder();
            rcb.setClassLoaders(new ClassLoader[] { getClass().getClassLoader() });
            rcb.addUrls(ClasspathHelper.forClassLoader(getClass().getClassLoader()));
            Reflections reflections = new Reflections(rcb);
            handlerClasses = reflections.getSubTypesOf(TransportHandler.class);
        }

        //        Collection<TransportHandler> result = new ArrayList<>();
        //        for (Class<? extends TransportHandler> handler : handlerClasses) {
        //            result.add((TransportHandler)injector.get(handler));
        //        }
        return new TransportHandlerSet(handlerClasses);
    }

    private FormatHandlerSet buildFormatHandlerSet() {
        Collection<Class<? extends FormatHandler>> handlerClasses;
        if (mc.workerTransportHandlers != null && !mc.workerTransportHandlers.isEmpty()) {
            handlerClasses = mc.workerFormatHandlers;
        } else {
            ConfigurationBuilder rcb = new ConfigurationBuilder();
            rcb.setClassLoaders(new ClassLoader[] { getClass().getClassLoader() });
            rcb.addUrls(ClasspathHelper.forClassLoader(getClass().getClassLoader()));
            Reflections reflections = new Reflections(rcb);
            handlerClasses = reflections.getSubTypesOf(FormatHandler.class);
        }

        Collection<FormatHandler> result = new ArrayList<>();
        for (Class<? extends FormatHandler> handler : handlerClasses) {
            result.add((FormatHandler) injector.get(handler));
        }
        return new FormatHandlerSet(result);
    }

    private void loadSeeds() throws IOException {
        Collection<File> seedFiles = new ArrayList<>();
        if (mc.conductorSeedFiles != null && !mc.conductorSeedFiles.isEmpty()) {
            seedFiles = mc.conductorSeedFiles;
        } else {
            seedFiles.add(new File(mc.getHomeDirectory(), "seeds.txt"));
        }
        if (seedFiles.isEmpty()) {
            return;
        }

        final HostService hostService = injector.get(HostService.class);
        final CatalogueService catalogueService = injector.get(CatalogueService.class);

        for (File seedFile : seedFiles) {
            if (!seedFile.exists()) {
                LOGGER.warn("Seed file {} does not exist", seedFile);
                continue;
            }

            Iterator<String> it = FileUtils.lineIterator(seedFile);
            while (it.hasNext()) {
                String line = it.next();
                if (line == null || line.trim().isEmpty() || line.trim().startsWith("#")) {
                    continue;
                }

                line = line.trim();

                io.mola.galimatias.URL url;
                try {
                    url = io.mola.galimatias.URL.parse(line);
                } catch (GalimatiasParseException ex) {
                    continue;
                }
                String tld = InternetDomainName.from(url.host().toString()).topPrivateDomain().toString();

                Host host = hostService.findByDomain(tld);
                if (host == null) {
                    host = new Host();
                    host.host = tld;
                }

                CatalogueEntry seed = new CatalogueEntry();
                seed.location = line;
                seed.initialDepth = 0;
                seed.status = Status.PENDING;
                seed.host = host;
                if (catalogueService.findForLocation(seed.location) == null) {
                    LOGGER.info("Found seed: {}", line);
                    catalogueService.save(seed);
                }
            }
        }
    }

    //~ Entry point ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    public static void main(String[] args) throws Exception {
        new Main(args);
    }

    private static class ExtensionBinder implements Binder {
        private final Injector injector;

        public ExtensionBinder(Injector injector) {
            this.injector = injector;
        }

        @Override
        public <T> Binder bind(Type type, T instance) {
            injector.bind(type, instance);
            return this;
        }

        @Override
        public <T> Binder bind(Type type, Type instance) {
            injector.bind(type, instance);
            return this;
        }

        @Override
        public <T> Binder bindToProvider(Type type, Provider<T> provider) {
            injector.bindToProvider(type, provider);
            return this;
        }

        @Override
        public <T> Binder bindToProvider(Type type, Class<? extends Provider<T>> provider) {
            injector.bindToProvider(type, provider);
            return this;
        }

    }
}