org.graylog2.inputs.transports.UdpTransportTest.java Source code

Java tutorial

Introduction

Here is the source code for org.graylog2.inputs.transports.UdpTransportTest.java

Source

/**
 * This file is part of Graylog.
 *
 * Graylog is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Graylog is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Graylog.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.graylog2.inputs.transports;

import com.codahale.metrics.Gauge;
import com.github.joschi.jadconfig.util.Size;
import com.google.common.collect.ImmutableMap;
import com.google.common.primitives.Ints;
import com.google.common.util.concurrent.Uninterruptibles;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandler;
import io.netty.channel.ChannelOption;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.FixedRecvByteBufAllocator;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.logging.LogLevel;
import io.netty.handler.logging.LoggingHandler;
import io.netty.util.ReferenceCountUtil;
import org.apache.commons.lang3.SystemUtils;
import org.graylog2.inputs.transports.netty.EventLoopGroupFactory;
import org.graylog2.plugin.LocalMetricRegistry;
import org.graylog2.plugin.configuration.Configuration;
import org.graylog2.plugin.configuration.ConfigurationRequest;
import org.graylog2.plugin.inputs.MessageInput;
import org.graylog2.plugin.inputs.MisfireException;
import org.graylog2.plugin.inputs.transports.NettyTransport;
import org.graylog2.plugin.inputs.util.ThroughputCounter;
import org.graylog2.shared.SuppressForbidden;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.TimeUnit;

import static com.jayway.awaitility.Awaitility.await;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.Assume.assumeTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

public class UdpTransportTest {
    private static final String BIND_ADDRESS = "127.0.0.1";
    private static final int PORT = 0;
    private static final int RECV_BUFFER_SIZE = 1024;
    private static final Map<String, Object> CONFIG_SOURCE = ImmutableMap.of(NettyTransport.CK_BIND_ADDRESS,
            BIND_ADDRESS, NettyTransport.CK_PORT, PORT, NettyTransport.CK_RECV_BUFFER_SIZE, RECV_BUFFER_SIZE,
            NettyTransport.CK_NUMBER_WORKER_THREADS, 1);
    private static final Configuration CONFIGURATION = new Configuration(CONFIG_SOURCE);

    private final NettyTransportConfiguration nettyTransportConfiguration = new NettyTransportConfiguration("nio",
            "jdk", 1);
    private UdpTransport udpTransport;
    private EventLoopGroup eventLoopGroup;
    private EventLoopGroupFactory eventLoopGroupFactory;
    private ThroughputCounter throughputCounter;
    private LocalMetricRegistry localMetricRegistry;

    @Before
    @SuppressForbidden("Executors#newSingleThreadExecutor() is okay for tests")
    public void setUp() throws Exception {
        eventLoopGroupFactory = new EventLoopGroupFactory(nettyTransportConfiguration);
        localMetricRegistry = new LocalMetricRegistry();
        eventLoopGroup = eventLoopGroupFactory.create(1, localMetricRegistry, "test");
        throughputCounter = new ThroughputCounter(eventLoopGroup);
        udpTransport = new UdpTransport(CONFIGURATION, eventLoopGroupFactory, nettyTransportConfiguration,
                throughputCounter, localMetricRegistry);
    }

    @After
    public void tearDown() {
        eventLoopGroup.shutdownGracefully();
    }

    @Test
    public void transportReceivesDataSmallerThanRecvBufferSize() throws Exception {
        final CountingChannelUpstreamHandler handler = new CountingChannelUpstreamHandler();
        final UdpTransport transport = launchTransportForBootStrapTest(handler);
        await().atMost(5, TimeUnit.SECONDS).until(() -> transport.getLocalAddress() != null);
        final InetSocketAddress localAddress = (InetSocketAddress) transport.getLocalAddress();

        sendUdpDatagram(BIND_ADDRESS, localAddress.getPort(), 100);
        await().atMost(5, TimeUnit.SECONDS).until(() -> !handler.getBytesWritten().isEmpty());
        transport.stop();

        assertThat(handler.getBytesWritten()).containsOnly(100);
    }

    @Test
    public void transportReceivesDataExactlyRecvBufferSize() throws Exception {
        final CountingChannelUpstreamHandler handler = new CountingChannelUpstreamHandler();
        final UdpTransport transport = launchTransportForBootStrapTest(handler);
        await().atMost(5, TimeUnit.SECONDS).until(() -> transport.getLocalAddress() != null);
        final InetSocketAddress localAddress = (InetSocketAddress) transport.getLocalAddress();

        // This will be variable depending on the version of the IP protocol and the UDP packet size.
        final int udpOverhead = 16;
        final int maxPacketSize = RECV_BUFFER_SIZE - udpOverhead;

        sendUdpDatagram(BIND_ADDRESS, localAddress.getPort(), maxPacketSize);
        await().atMost(5, TimeUnit.SECONDS).until(() -> !handler.getBytesWritten().isEmpty());
        transport.stop();

        assertThat(handler.getBytesWritten()).containsOnly(maxPacketSize);
    }

    @Test
    public void transportDiscardsDataLargerRecvBufferSizeOnMacOsX() throws Exception {
        assumeTrue("Skipping test intended for MacOS X systems", SystemUtils.IS_OS_MAC_OSX);

        final CountingChannelUpstreamHandler handler = new CountingChannelUpstreamHandler();
        final UdpTransport transport = launchTransportForBootStrapTest(handler);
        await().atMost(5, TimeUnit.SECONDS).until(() -> transport.getLocalAddress() != null);
        final InetSocketAddress localAddress = (InetSocketAddress) transport.getLocalAddress();

        sendUdpDatagram(BIND_ADDRESS, localAddress.getPort(), 2 * RECV_BUFFER_SIZE);
        Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);

        transport.stop();

        assertThat(handler.getBytesWritten()).isEmpty();
    }

    @Test
    public void transportTruncatesDataLargerRecvBufferSizeOnLinux() throws Exception {
        assumeTrue("Skipping test intended for Linux systems", SystemUtils.IS_OS_LINUX);

        final CountingChannelUpstreamHandler handler = new CountingChannelUpstreamHandler();
        final UdpTransport transport = launchTransportForBootStrapTest(handler);
        await().atMost(5, TimeUnit.SECONDS).until(() -> transport.getLocalAddress() != null);
        final InetSocketAddress localAddress = (InetSocketAddress) transport.getLocalAddress();

        sendUdpDatagram(BIND_ADDRESS, localAddress.getPort(), 2 * RECV_BUFFER_SIZE);
        await().atMost(5, TimeUnit.SECONDS).until(() -> !handler.getBytesWritten().isEmpty());
        transport.stop();

        assertThat(handler.getBytesWritten()).containsExactly(RECV_BUFFER_SIZE);
    }

    @Test
    public void receiveBufferSizeIsDefaultSize() {
        assertThat(
                udpTransport.getBootstrap(mock(MessageInput.class)).config().options().get(ChannelOption.SO_RCVBUF))
                        .isEqualTo(RECV_BUFFER_SIZE);
    }

    @Test
    public void receiveBufferSizeIsNotLimited() {
        final int recvBufferSize = Ints.saturatedCast(Size.megabytes(1L).toBytes());
        ImmutableMap<String, Object> source = ImmutableMap.of(NettyTransport.CK_BIND_ADDRESS, BIND_ADDRESS,
                NettyTransport.CK_PORT, PORT, NettyTransport.CK_RECV_BUFFER_SIZE, recvBufferSize);
        Configuration config = new Configuration(source);
        UdpTransport udpTransport = new UdpTransport(config, eventLoopGroupFactory, nettyTransportConfiguration,
                throughputCounter, new LocalMetricRegistry());

        assertThat(
                udpTransport.getBootstrap(mock(MessageInput.class)).config().options().get(ChannelOption.SO_RCVBUF))
                        .isEqualTo(recvBufferSize);
    }

    @Test
    public void receiveBufferSizePredictorIsUsingDefaultSize() {
        FixedRecvByteBufAllocator recvByteBufAllocator = (FixedRecvByteBufAllocator) udpTransport
                .getBootstrap(mock(MessageInput.class)).config().options().get(ChannelOption.RCVBUF_ALLOCATOR);
        assertThat(recvByteBufAllocator.newHandle().guess()).isEqualTo(RECV_BUFFER_SIZE);
    }

    @Test
    public void getMetricSetReturnsLocalMetricRegistry() {
        assertThat(udpTransport.getMetricSet()).isSameAs(localMetricRegistry);
    }

    @Test
    public void testDefaultReceiveBufferSize() {
        final UdpTransport.Config config = new UdpTransport.Config();
        final ConfigurationRequest requestedConfiguration = config.getRequestedConfiguration();

        assertThat(requestedConfiguration.getField(NettyTransport.CK_RECV_BUFFER_SIZE).getDefaultValue())
                .isEqualTo(262144);
    }

    @Test
    public void testTrafficCounter() throws Exception {
        final CountingChannelUpstreamHandler handler = new CountingChannelUpstreamHandler();
        final UdpTransport transport = launchTransportForBootStrapTest(handler);
        try {
            await().atMost(5, TimeUnit.SECONDS).until(() -> transport.getLocalAddress() != null);
            final InetSocketAddress localAddress = (InetSocketAddress) transport.getLocalAddress();

            sendUdpDatagram(BIND_ADDRESS, localAddress.getPort(), 512);
            await().atMost(5, TimeUnit.SECONDS).until(() -> handler.getBytesWritten().size() == 1);
            assertThat(handler.getBytesWritten()).containsExactly(512);

            sendUdpDatagram(BIND_ADDRESS, localAddress.getPort(), 512);
            await().atMost(5, TimeUnit.SECONDS).until(() -> handler.getBytesWritten().size() == 2);
            assertThat(handler.getBytesWritten()).containsExactly(512, 512);

        } finally {
            transport.stop();
        }

        final Map<String, Gauge<Long>> gauges = throughputCounter.gauges();
        assertThat(gauges.get(ThroughputCounter.READ_BYTES_TOTAL).getValue()).isEqualTo(1024L);
    }

    private UdpTransport launchTransportForBootStrapTest(final ChannelInboundHandler channelHandler)
            throws MisfireException {
        final UdpTransport transport = new UdpTransport(CONFIGURATION, eventLoopGroupFactory,
                nettyTransportConfiguration, throughputCounter, new LocalMetricRegistry()) {
            @Override
            protected LinkedHashMap<String, Callable<? extends ChannelHandler>> getChannelHandlers(
                    MessageInput input) {
                final LinkedHashMap<String, Callable<? extends ChannelHandler>> handlers = new LinkedHashMap<>();
                handlers.put("logging", () -> new LoggingHandler(LogLevel.INFO));
                handlers.put("counter", () -> channelHandler);
                handlers.putAll(super.getChannelHandlers(input));
                return handlers;
            }
        };

        final MessageInput messageInput = mock(MessageInput.class);
        when(messageInput.getId()).thenReturn("TEST");
        when(messageInput.getName()).thenReturn("TEST");

        transport.launch(messageInput);

        return transport;
    }

    private void sendUdpDatagram(String hostname, int port, int size) throws IOException {
        final InetAddress address = InetAddress.getByName(hostname);
        final byte[] data = new byte[size];
        final DatagramPacket packet = new DatagramPacket(data, data.length, address, port);

        try (DatagramSocket toSocket = new DatagramSocket()) {
            toSocket.send(packet);
        }
    }

    public static class CountingChannelUpstreamHandler
            extends SimpleChannelInboundHandler<io.netty.channel.socket.DatagramPacket> {
        private final List<Integer> bytesWritten = new CopyOnWriteArrayList<>();

        @Override
        protected void channelRead0(ChannelHandlerContext ctx, io.netty.channel.socket.DatagramPacket msg)
                throws Exception {
            bytesWritten.add(msg.content().readableBytes());
            ctx.fireChannelRead(ReferenceCountUtil.retain(msg));
        }

        List<Integer> getBytesWritten() {
            return bytesWritten;
        }
    }
}