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

Introduction

Here is the source code for com.github.brandtg.switchboard.LogReceiver.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.bootstrap.ServerBootstrap;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufInputStream;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;

public class LogReceiver {
    private static final Logger LOG = LoggerFactory.getLogger(LogReceiver.class);
    private static final int MAX_FRAME_LENGTH = 1024 * 1024;
    private static final int LENGTH_FIELD_OFFSET = 0;
    private static final int LENGTH_FIELD_LENGTH = 4;
    private static final int LENGTH_ADJUSTMENT = 0;
    private static final int INITIAL_BYTES_TO_STRIP = 4;

    private final InetSocketAddress address;
    private final AtomicBoolean isShutdown;
    private final Set<Object> listeners;
    private final ServerBootstrap serverBootstrap;

    private ChannelFuture bootstrapFuture;

    /**
     * A server that listens for log regions and pipes them to an output stream.
     *
     * @param address
     *  The socket address on which to listen
     * @param eventExecutors
     *  The Netty executor service to use for incoming traffic
     * @param outputStream
     *  The output stream to which all data should be piped
     */
    public LogReceiver(InetSocketAddress address, EventLoopGroup eventExecutors, final OutputStream outputStream) {
        this.address = address;
        this.isShutdown = new AtomicBoolean(true);
        this.listeners = new HashSet<Object>();
        this.serverBootstrap = new ServerBootstrap().group(eventExecutors).channel(NioServerSocketChannel.class)
                .childHandler(new ChannelInitializer<SocketChannel>() {
                    @Override
                    protected void initChannel(SocketChannel ch) throws Exception {
                        ch.pipeline()
                                .addLast(new LengthFieldBasedFrameDecoder(MAX_FRAME_LENGTH, LENGTH_FIELD_OFFSET,
                                        LENGTH_FIELD_LENGTH, LENGTH_ADJUSTMENT, INITIAL_BYTES_TO_STRIP));
                        ch.pipeline().addLast(new LogMessageHandler(outputStream, listeners));
                    }
                });
    }

    /**
     * Start listening for log traffic.
     */
    public void start() throws Exception {
        if (isShutdown.getAndSet(false)) {
            bootstrapFuture = serverBootstrap.bind(address).sync();
        }
    }

    /**
     * Stop listening for log traffic.
     */
    public void shutdown() throws Exception {
        if (isShutdown.getAndSet(true)) {
            if (bootstrapFuture != null) {
                bootstrapFuture.channel().close();
            }
        }
    }

    /**
     * Returns the local socket address to which log data can be sent.
     */
    public SocketAddress getLocalAddress() throws Exception {
        if (bootstrapFuture == null) {
            return null;
        }
        return bootstrapFuture.channel().localAddress();
    }

    /**
     * Registers a listener, on which {@link Object#notify()} will be called when data arrives.
     *
     * <p>
     *   n.b. Other object should not synchronize on listener to avoid deadlocks.
     * </p>
     */
    public void registerListener(Object listener) {
        synchronized (listeners) {
            listeners.add(listener);
        }
    }

    private static class LogMessageHandler extends SimpleChannelInboundHandler<ByteBuf> {
        private final OutputStream outputStream;
        private final Set<Object> listeners;

        LogMessageHandler(OutputStream outputStream, Set<Object> listeners) {
            this.outputStream = outputStream;
            this.listeners = listeners;
        }

        @Override
        protected void channelRead0(ChannelHandlerContext ctx, ByteBuf msg) throws Exception {
            synchronized (listeners) {
                int readableBytes = msg.readableBytes();
                IOUtils.copy(new ByteBufInputStream(msg), outputStream);
                outputStream.flush();
                LOG.debug("Wrote {} to output stream", readableBytes);

                for (Object listener : listeners) {
                    synchronized (listener) {
                        LOG.debug("Notifying {}", listener);
                        listener.notify();
                    }
                }
            }
        }
    }

    public static void main(String[] args) throws Exception {
        LogReceiver logReceiver = new LogReceiver(new InetSocketAddress(2000), new NioEventLoopGroup(), System.out);
        logReceiver.start();
    }
}