com.github.brandtg.switchboard.FileLogAggregator.java Source code

Java tutorial

Introduction

Here is the source code for com.github.brandtg.switchboard.FileLogAggregator.java

Source

/**
 * Copyright (C) 2015 Greg Brandt (brandt.greg@gmail.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 com.github.brandtg.switchboard;

import io.netty.channel.EventLoopGroup;
import io.netty.channel.nio.NioEventLoopGroup;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.GnuParser;
import org.apache.commons.cli.HelpFormatter;
import org.apache.commons.cli.Options;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.FileOutputStream;
import java.io.InputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.PipedInputStream;
import java.io.PipedOutputStream;
import java.net.InetSocketAddress;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Executors;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicBoolean;

public class FileLogAggregator {
    private static final Logger LOG = LoggerFactory.getLogger(FileLogAggregator.class);

    private final char separator;
    private final OutputStream outputStream;
    private final AtomicBoolean isShutdown;
    private final EventLoopGroup eventExecutors;
    private final Map<InetSocketAddress, LogReceiver> logReceivers;
    private final Map<InetSocketAddress, InputStream> inputStreams;
    private final Map<InetSocketAddress, StringBuffer> stringBuffers;
    private final Object listener;
    private final ExecutorService muxExecutor;

    /**
     * An agent that aggregates logs from multiple sources and multiplexes them.
     *
     * @param sources
     *  A set of source switchboard servers from which to pull logs
     * @param separator
     *  The line delimiter (after this is reached, lines will be output to outputStream)
     * @param outputStream
     *  OutputStream to which all multiplexed logs are piped
     */
    public FileLogAggregator(Set<InetSocketAddress> sources, char separator, OutputStream outputStream)
            throws IOException {
        this.separator = separator;
        this.outputStream = outputStream;
        this.isShutdown = new AtomicBoolean(true);
        this.eventExecutors = new NioEventLoopGroup();
        this.logReceivers = new HashMap<>();
        this.inputStreams = new HashMap<>();
        this.listener = new Object();
        this.muxExecutor = Executors.newSingleThreadExecutor();
        this.stringBuffers = new HashMap<>();

        for (InetSocketAddress source : sources) {
            PipedOutputStream pos = new PipedOutputStream();
            PipedInputStream pis = new PipedInputStream(pos);
            LogReceiver logReceiver = new LogReceiver(new InetSocketAddress(0), eventExecutors, pos);
            logReceiver.registerListener(listener);
            logReceivers.put(source, logReceiver);
            inputStreams.put(source, pis);
            stringBuffers.put(source, new StringBuffer());
        }
    }

    /**
     * Starts pulling and receiving logs from sources.
     */
    public void start() throws Exception {
        if (isShutdown.getAndSet(false)) {
            muxExecutor.submit(new LogMultiplexer());
            for (Map.Entry<InetSocketAddress, LogReceiver> entry : logReceivers.entrySet()) {
                entry.getValue().start();
                LOG.info("Started log receiver for {} on {}", entry.getKey(), entry.getValue().getLocalAddress());
            }
        }
    }

    /**
     * Stops pulling and receiving logs from sources.
     */
    public void stop() throws Exception {
        if (!isShutdown.getAndSet(true)) {
            for (Map.Entry<InetSocketAddress, LogReceiver> entry : logReceivers.entrySet()) {
                entry.getValue().shutdown();
                LOG.info("Stopped log receiver for {} on local {}", entry.getKey(),
                        entry.getValue().getLocalAddress());
            }
            eventExecutors.shutdownGracefully();
        }
    }

    private class LogMultiplexer implements Runnable {
        @Override
        public void run() {
            while (!isShutdown.get()) {
                try {
                    // Wait until some data has been received
                    synchronized (listener) {
                        listener.wait();
                    }

                    // Load data into string buffers and print lines
                    for (Map.Entry<InetSocketAddress, InputStream> entry : inputStreams.entrySet()) {
                        InputStream is = entry.getValue();
                        StringBuffer buf = stringBuffers.get(entry.getKey());
                        int count = is.available();
                        for (int i = 0; i < count; i++) {
                            char nextChar = (char) is.read();
                            if (nextChar == separator) {
                                byte[] line = buf.toString().getBytes();
                                buf.setLength(0);
                                outputStream.write(line);
                                outputStream.write(separator);
                            } else {
                                buf.append(nextChar);
                            }
                        }
                    }
                } catch (Exception e) {
                    LOG.error("Error in Multiplexer", e);
                }
            }
        }
    }

    /** Main. */
    public static void main(String[] args) throws Exception {
        Options opts = new Options();
        opts.addOption("h", "help", false, "Prints help message");
        opts.addOption("f", "file", true, "File to output aggregated logs to");
        opts.addOption("s", "separator", true, "Line separator in log");
        CommandLine cli = new GnuParser().parse(opts, args);

        if (cli.getArgs().length == 0 || cli.hasOption("help")) {
            new HelpFormatter().printHelp("usage: [opts] sourceHost:port ...", opts);
            System.exit(1);
        }

        // Parse sources
        Set<InetSocketAddress> sources = new HashSet<>();
        for (int i = 0; i < cli.getArgs().length; i++) {
            String[] tokens = cli.getArgs()[i].split(":");
            sources.add(new InetSocketAddress(tokens[0], Integer.valueOf(tokens[1])));
        }

        // Parse output stream
        OutputStream outputStream;
        if (cli.hasOption("file")) {
            outputStream = new FileOutputStream(cli.getOptionValue("file"));
        } else {
            outputStream = System.out;
        }

        // Separator
        String separator = cli.getOptionValue("separator", "\n");
        if (separator.length() != 1) {
            throw new IllegalArgumentException("Separator must only be 1 character");
        }

        final FileLogAggregator fileLogAggregator = new FileLogAggregator(sources, separator.charAt(0),
                outputStream);

        Runtime.getRuntime().addShutdownHook(new Thread() {
            @Override
            public void run() {
                try {
                    fileLogAggregator.stop();
                } catch (Exception e) {
                    LOG.error("Error when stopping log aggregator", e);
                }
            }
        });

        fileLogAggregator.start();
    }
}