Java tutorial
/** * 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.github.joschi.jadconfig.util.Size; import com.google.common.collect.ImmutableMap; import com.google.common.primitives.Ints; import com.google.common.util.concurrent.Callables; import com.google.common.util.concurrent.Uninterruptibles; import org.apache.commons.lang3.SystemUtils; 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.jboss.netty.buffer.ChannelBuffer; import org.jboss.netty.channel.ChannelHandler; import org.jboss.netty.channel.ChannelHandlerContext; import org.jboss.netty.channel.FixedReceiveBufferSizePredictorFactory; import org.jboss.netty.channel.MessageEvent; import org.jboss.netty.channel.ReceiveBufferSizePredictorFactory; import org.jboss.netty.channel.SimpleChannelUpstreamHandler; import org.jboss.netty.channel.UpstreamMessageEvent; import org.jboss.netty.util.HashedWheelTimer; 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.ArrayList; import java.util.LinkedHashMap; import java.util.List; import java.util.concurrent.Callable; 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 ImmutableMap<String, Object> CONFIG_SOURCE = ImmutableMap.<String, Object>of( NettyTransport.CK_BIND_ADDRESS, BIND_ADDRESS, NettyTransport.CK_PORT, PORT, NettyTransport.CK_RECV_BUFFER_SIZE, RECV_BUFFER_SIZE); private static final Configuration CONFIGURATION = new Configuration(CONFIG_SOURCE); private UdpTransport udpTransport; private ThroughputCounter throughputCounter; private LocalMetricRegistry localMetricRegistry; @Before public void setUp() throws Exception { throughputCounter = new ThroughputCounter(new HashedWheelTimer()); localMetricRegistry = new LocalMetricRegistry(); udpTransport = new UdpTransport(CONFIGURATION, throughputCounter, localMetricRegistry); } @Test public void transportReceivesDataSmallerThanRecvBufferSize() throws Exception { final CountingChannelUpstreamHandler handler = new CountingChannelUpstreamHandler(); final UdpTransport transport = launchTransportForBootStrapTest(handler); final InetSocketAddress localAddress = (InetSocketAddress) transport.getLocalAddress(); sendUdpDatagram(BIND_ADDRESS, localAddress.getPort(), 100); await().atMost(5, TimeUnit.SECONDS).until(new Callable<Boolean>() { @Override public Boolean call() throws Exception { return !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); 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(new Callable<Boolean>() { @Override public Boolean call() throws Exception { return !handler.getBytesWritten().isEmpty(); } }); transport.stop(); assertThat(handler.getBytesWritten()).containsOnly(maxPacketSize); } @Test public void transportDiscardsDataLargerRecvBufferSizeOnMacOsX() throws Exception { assumeTrue(SystemUtils.IS_OS_MAC_OSX); final CountingChannelUpstreamHandler handler = new CountingChannelUpstreamHandler(); final UdpTransport transport = launchTransportForBootStrapTest(handler); 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(SystemUtils.IS_OS_LINUX); final CountingChannelUpstreamHandler handler = new CountingChannelUpstreamHandler(); final UdpTransport transport = launchTransportForBootStrapTest(handler); final InetSocketAddress localAddress = (InetSocketAddress) transport.getLocalAddress(); sendUdpDatagram(BIND_ADDRESS, localAddress.getPort(), 2 * RECV_BUFFER_SIZE); await().atMost(5, TimeUnit.SECONDS).until(new Callable<Boolean>() { @Override public Boolean call() throws Exception { return !handler.getBytesWritten().isEmpty(); } }); transport.stop(); assertThat(handler.getBytesWritten()).containsExactly(RECV_BUFFER_SIZE); } private UdpTransport launchTransportForBootStrapTest(final ChannelHandler channelHandler) throws MisfireException { final UdpTransport transport = new UdpTransport(CONFIGURATION, throughputCounter, new LocalMetricRegistry()) { @Override protected LinkedHashMap<String, Callable<? extends ChannelHandler>> getBaseChannelHandlers( MessageInput input) { final LinkedHashMap<String, Callable<? extends ChannelHandler>> handlers = new LinkedHashMap<>(); handlers.put("counter", Callables.returning(channelHandler)); handlers.putAll(super.getFinalChannelHandlers(input)); return handlers; } }; final MessageInput messageInput = mock(MessageInput.class); when(messageInput.getId()).thenReturn("TEST"); when(messageInput.getName()).thenReturn("TEST"); transport.launch(messageInput); return transport; } @Test public void receiveBufferSizeIsDefaultSize() throws Exception { assertThat(udpTransport.getBootstrap().getOption("receiveBufferSize")).isEqualTo(RECV_BUFFER_SIZE); } @Test public void receiveBufferSizeIsNotLimited() throws Exception { final int recvBufferSize = Ints.saturatedCast(Size.megabytes(1L).toBytes()); ImmutableMap<String, Object> source = ImmutableMap.<String, Object>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, throughputCounter, new LocalMetricRegistry()); assertThat(udpTransport.getBootstrap().getOption("receiveBufferSize")).isEqualTo(recvBufferSize); } @Test public void receiveBufferSizePredictorIsUsingDefaultSize() throws Exception { ReceiveBufferSizePredictorFactory receiveBufferSizePredictorFactory = (ReceiveBufferSizePredictorFactory) udpTransport .getBootstrap().getOption("receiveBufferSizePredictorFactory"); assertThat(receiveBufferSizePredictorFactory).isInstanceOf(FixedReceiveBufferSizePredictorFactory.class); assertThat(receiveBufferSizePredictorFactory.getPredictor().nextReceiveBufferSize()) .isEqualTo(RECV_BUFFER_SIZE); } @Test public void getMetricSetReturnsLocalMetricRegistry() { assertThat(udpTransport.getMetricSet()).isSameAs(localMetricRegistry); } @Test public void testDefaultReceiveBufferSize() throws Exception { final UdpTransport.Config config = new UdpTransport.Config(); final ConfigurationRequest requestedConfiguration = config.getRequestedConfiguration(); assertThat(requestedConfiguration.getField(NettyTransport.CK_RECV_BUFFER_SIZE).getDefaultValue()) .isEqualTo(262144); } 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); DatagramSocket toSocket = null; try { toSocket = new DatagramSocket(); toSocket.send(packet); } finally { if (toSocket != null) { toSocket.close(); } } } public static class CountingChannelUpstreamHandler extends SimpleChannelUpstreamHandler { private final List<Integer> bytesWritten = new ArrayList<>(); @Override public void messageReceived(ChannelHandlerContext ctx, MessageEvent e) throws Exception { if (e instanceof UpstreamMessageEvent) { ChannelBuffer buffer = (ChannelBuffer) e.getMessage(); try { synchronized (bytesWritten) { bytesWritten.add(buffer.readableBytes()); } } finally { super.messageReceived(ctx, e); } } } public List<Integer> getBytesWritten() { return bytesWritten; } } }