org.graylog2.inputs.codecs.GelfChunkAggregatorTest.java Source code

Java tutorial

Introduction

Here is the source code for org.graylog2.inputs.codecs.GelfChunkAggregatorTest.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.codecs;

import com.codahale.metrics.Metric;
import com.codahale.metrics.MetricFilter;
import com.codahale.metrics.MetricRegistry;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import org.graylog2.plugin.InstantMillisProvider;
import org.graylog2.plugin.inputs.codecs.CodecAggregator;
import org.joda.time.DateTime;
import org.joda.time.DateTimeUtils;
import org.joda.time.DateTimeZone;
import org.joda.time.Period;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.mockito.junit.MockitoJUnit;
import org.mockito.junit.MockitoRule;

import java.util.concurrent.ConcurrentSkipListSet;
import java.util.concurrent.ScheduledThreadPoolExecutor;

import static org.assertj.core.api.Assertions.assertThat;
import static org.graylog2.inputs.codecs.GelfChunkAggregator.CHUNK_COUNTER;
import static org.graylog2.inputs.codecs.GelfChunkAggregator.COMPLETE_MESSAGES;
import static org.graylog2.inputs.codecs.GelfChunkAggregator.ChunkEntry;
import static org.graylog2.inputs.codecs.GelfChunkAggregator.DUPLICATE_CHUNKS;
import static org.graylog2.inputs.codecs.GelfChunkAggregator.EXPIRED_CHUNKS;
import static org.graylog2.inputs.codecs.GelfChunkAggregator.EXPIRED_MESSAGES;
import static org.graylog2.inputs.codecs.GelfChunkAggregator.WAITING_MESSAGES;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;

public class GelfChunkAggregatorTest {
    private static final byte[] CHUNK_MAGIC_BYTES = new byte[] { 0x1e, 0x0f };

    @Rule
    public final MockitoRule mockitoRule = MockitoJUnit.rule();

    private ScheduledThreadPoolExecutor poolExecutor;
    private GelfChunkAggregator aggregator;
    private MetricRegistry metricRegistry;

    @Before
    public void before() {
        poolExecutor = new ScheduledThreadPoolExecutor(1);
        metricRegistry = new MetricRegistry();
        aggregator = new GelfChunkAggregator(poolExecutor, metricRegistry);
    }

    @After
    public void after() {
        poolExecutor.shutdown();
        DateTimeUtils.setCurrentMillisSystem();
    }

    @Test
    public void addSingleChunk() {
        final ByteBuf[] singleChunk = createChunkedMessage(512, 1024);

        final CodecAggregator.Result result = aggregator.addChunk(singleChunk[0]);

        assertNotNull("message should be complete", result.getMessage());

        assertEquals(1, counterValueNamed(metricRegistry, COMPLETE_MESSAGES));
        assertEquals(1, counterValueNamed(metricRegistry, CHUNK_COUNTER));
        assertEquals(0, counterValueNamed(metricRegistry, WAITING_MESSAGES));
        assertEquals(0, counterValueNamed(metricRegistry, EXPIRED_CHUNKS));
        assertEquals(0, counterValueNamed(metricRegistry, EXPIRED_MESSAGES));
        assertEquals(0, counterValueNamed(metricRegistry, DUPLICATE_CHUNKS));
    }

    @Test
    public void manyChunks() {
        final ByteBuf[] chunks = createChunkedMessage(4096 + 512, 1024); // creates 5 chunks

        int i = 0;
        for (final ByteBuf chunk : chunks) {
            i++;
            final CodecAggregator.Result result = aggregator.addChunk(chunk);
            assertTrue(result.isValid());
            if (i == 5) {
                assertNotNull("message should've been assembled from chunks", result.getMessage());

                assertEquals(1, counterValueNamed(metricRegistry, COMPLETE_MESSAGES));
                assertEquals(5, counterValueNamed(metricRegistry, CHUNK_COUNTER));
                assertEquals(0, counterValueNamed(metricRegistry, WAITING_MESSAGES));
                assertEquals(0, counterValueNamed(metricRegistry, EXPIRED_CHUNKS));
                assertEquals(0, counterValueNamed(metricRegistry, EXPIRED_MESSAGES));
                assertEquals(0, counterValueNamed(metricRegistry, DUPLICATE_CHUNKS));

            } else {
                assertNull("chunks not complete", result.getMessage());

                assertEquals("message not complete yet", 0, counterValueNamed(metricRegistry, COMPLETE_MESSAGES));
                assertEquals(i, counterValueNamed(metricRegistry, CHUNK_COUNTER));
                assertEquals("one message waiting", 1, counterValueNamed(metricRegistry, WAITING_MESSAGES));
                assertEquals(0, counterValueNamed(metricRegistry, EXPIRED_CHUNKS));
                assertEquals(0, counterValueNamed(metricRegistry, EXPIRED_MESSAGES));
                assertEquals(0, counterValueNamed(metricRegistry, DUPLICATE_CHUNKS));
            }
        }
    }

    @Test
    public void tooManyChunks() {
        final ByteBuf[] chunks = createChunkedMessage(129 * 1024, 1024);
        int i = 1;
        for (final ByteBuf chunk : chunks) {
            final CodecAggregator.Result result = aggregator.addChunk(chunk);
            if (i == 129) {
                assertFalse("Message invalidated (chunk #" + i + ")", result.isValid());
                assertNull("Message discarded (chunk #" + i + ")", result.getMessage());

            } else {
                assertTrue("Incomplete message valid (chunk #" + i + ")", result.isValid());
                assertNull("Message not complete (chunk #" + i + ")", result.getMessage());
            }
            i++;
        }
    }

    @Test
    public void missingChunk() {
        final DateTime initialTime = new DateTime(2014, 1, 1, 1, 59, 59, 0, DateTimeZone.UTC);
        final InstantMillisProvider clock = new InstantMillisProvider(initialTime);
        DateTimeUtils.setCurrentMillisProvider(clock);

        // we don't want the clean up task to run automatically
        poolExecutor = mock(ScheduledThreadPoolExecutor.class);
        final MetricRegistry metricRegistry = new MetricRegistry();
        aggregator = new GelfChunkAggregator(poolExecutor, metricRegistry);
        final GelfChunkAggregator.ChunkEvictionTask evictionTask = aggregator.new ChunkEvictionTask();

        final ByteBuf[] chunks = createChunkedMessage(4096 + 512, 1024); // creates 5 chunks

        int i = 0;
        for (final ByteBuf chunk : chunks) {
            final CodecAggregator.Result result;
            // skip first chunk
            if (i++ == 0) {
                continue;
            }
            result = aggregator.addChunk(chunk);
            assertTrue(result.isValid());
            assertNull("chunks not complete", result.getMessage());
        }
        // move clock forward enough to evict all of the chunks
        clock.tick(Period.seconds(10));

        evictionTask.run();

        final CodecAggregator.Result result = aggregator.addChunk(chunks[0]);

        assertNull("message should not be complete because chunks were evicted already", result.getMessage());
        assertTrue(result.isValid());

        // we send all chunks but the last one comes too late
        assertEquals("no message is complete", 0, counterValueNamed(metricRegistry, COMPLETE_MESSAGES));
        assertEquals("received 5 chunks", 5, counterValueNamed(metricRegistry, CHUNK_COUNTER));
        assertEquals("last chunk creates another waiting message", 1,
                counterValueNamed(metricRegistry, WAITING_MESSAGES));
        assertEquals("4 chunks expired", 4, counterValueNamed(metricRegistry, EXPIRED_CHUNKS));
        assertEquals("one message expired", 1, counterValueNamed(metricRegistry, EXPIRED_MESSAGES));
        assertEquals("no duplicate chunks", 0, counterValueNamed(metricRegistry, DUPLICATE_CHUNKS));

        // reset clock for other tests
        DateTimeUtils.setCurrentMillisSystem();
    }

    @Test
    public void outOfOrderChunks() {
        final ByteBuf[] chunks = createChunkedMessage(4096 + 512, 1024); // creates 5 chunks
        CodecAggregator.Result result = null;
        for (int i = chunks.length - 1; i >= 0; i--) {
            result = aggregator.addChunk(chunks[i]);
            if (i != 0) {
                assertNull("message still incomplete", result.getMessage());
            }
        }
        assertNotNull(result);
        assertNotNull("first chunk should've completed the message", result.getMessage());
        assertEquals(1, counterValueNamed(metricRegistry, COMPLETE_MESSAGES));
        assertEquals(5, counterValueNamed(metricRegistry, CHUNK_COUNTER));
        assertEquals(0, counterValueNamed(metricRegistry, WAITING_MESSAGES));
        assertEquals(0, counterValueNamed(metricRegistry, EXPIRED_CHUNKS));
        assertEquals(0, counterValueNamed(metricRegistry, EXPIRED_MESSAGES));
        assertEquals(0, counterValueNamed(metricRegistry, DUPLICATE_CHUNKS));
    }

    @Test
    public void differentIdsDoNotInterfere() {
        final ByteBuf[] msg1 = createChunkedMessage(4096 + 1, 1024, generateMessageId(1));// 5 chunks;
        final ByteBuf[] msg2 = createChunkedMessage(4096 + 1, 1024, generateMessageId(2));// 5 chunks;

        CodecAggregator.Result result1 = null;
        CodecAggregator.Result result2 = null;
        for (int i = 0; i < msg1.length; i++) {
            result1 = aggregator.addChunk(msg1[i]);
            if (i > 0) {
                result2 = aggregator.addChunk(msg2[i]);
            }
        }
        assertNotNull(result1);
        assertNotNull(result2);
        assertNotNull("message 1 should be complete", result1.getMessage());
        assertNull("message 2 should not be complete", result2.getMessage());
        // only one is complete, we sent 9 chunks
        assertEquals(1, counterValueNamed(metricRegistry, COMPLETE_MESSAGES));
        assertEquals(9, counterValueNamed(metricRegistry, CHUNK_COUNTER));
        assertEquals(1, counterValueNamed(metricRegistry, WAITING_MESSAGES));
        assertEquals(0, counterValueNamed(metricRegistry, EXPIRED_CHUNKS));
        assertEquals(0, counterValueNamed(metricRegistry, EXPIRED_MESSAGES));
        assertEquals(0, counterValueNamed(metricRegistry, DUPLICATE_CHUNKS));
    }

    @Test
    public void duplicateChunk() {
        final byte[] messageId1 = generateMessageId(1);
        final byte[] messageId2 = generateMessageId(2);
        final ByteBuf chunk1 = createChunk(messageId1, (byte) 0, (byte) 2, new byte[16]);
        final ByteBuf chunk2 = createChunk(messageId1, (byte) 0, (byte) 2, new byte[16]);
        final ByteBuf chunk3 = createChunk(messageId2, (byte) 0, (byte) 2, new byte[16]);
        final ByteBuf chunk4 = createChunk(messageId1, (byte) 1, (byte) 2, new byte[16]);
        final ByteBuf chunk5 = createChunk(messageId2, (byte) 1, (byte) 2, new byte[16]);

        assertNull("message should not be complete", aggregator.addChunk(chunk1).getMessage());
        assertNull("message should not be complete", aggregator.addChunk(chunk2).getMessage());
        assertNull("message should not be complete", aggregator.addChunk(chunk3).getMessage());
        assertNotNull("message 1 should be complete", aggregator.addChunk(chunk4).getMessage());
        assertNotNull("message 2 should be complete", aggregator.addChunk(chunk5).getMessage());

        assertEquals(2, counterValueNamed(metricRegistry, COMPLETE_MESSAGES));
        assertEquals(5, counterValueNamed(metricRegistry, CHUNK_COUNTER));
        assertEquals(0, counterValueNamed(metricRegistry, WAITING_MESSAGES));
        assertEquals(0, counterValueNamed(metricRegistry, EXPIRED_CHUNKS));
        assertEquals(0, counterValueNamed(metricRegistry, EXPIRED_MESSAGES));
        assertEquals(1, counterValueNamed(metricRegistry, DUPLICATE_CHUNKS));
    }

    @Test
    public void testChunkEntryCompareTo() throws Exception {
        // Test if the ChunkEntry#compareTo() method can handle ChunkEntry objects which have the same timestamp.
        // See: https://github.com/Graylog2/graylog2-server/issues/1462

        final ConcurrentSkipListSet<GelfChunkAggregator.ChunkEntry> sortedEvictionSet = new ConcurrentSkipListSet<>();
        final long currentTime = System.currentTimeMillis();

        for (int i = 0; i < 10; i++) {
            sortedEvictionSet.add(new GelfChunkAggregator.ChunkEntry(1, currentTime, "a" + i));
        }

        final int size = sortedEvictionSet.size();

        for (int i = 0; i < size; i++) {
            sortedEvictionSet.remove(sortedEvictionSet.first());
        }

        assertTrue("eviction set should be empty", sortedEvictionSet.isEmpty());
    }

    @Test
    public void testChunkEntryEquals() throws Exception {
        final GelfChunkAggregator.ChunkEntry entry = new ChunkEntry(1, 0L, "id");

        assertThat(entry).isEqualTo(new ChunkEntry(1, 0L, "id"));
        assertThat(entry).isEqualTo(new ChunkEntry(2, 0L, "id"));
        assertThat(entry).isNotEqualTo(new ChunkEntry(1, 1L, "id"));
        assertThat(entry).isNotEqualTo(new ChunkEntry(1, 0L, "foo"));
    }

    @Test
    public void testChunkEntryHashCode() throws Exception {
        final GelfChunkAggregator.ChunkEntry entry = new ChunkEntry(1, 0L, "id");

        assertThat(entry.hashCode()).isEqualTo(new ChunkEntry(1, 0L, "id").hashCode());
        assertThat(entry.hashCode()).isEqualTo(new ChunkEntry(2, 0L, "id").hashCode());
        assertThat(entry.hashCode()).isNotEqualTo(new ChunkEntry(1, 1L, "id").hashCode());
        assertThat(entry.hashCode()).isNotEqualTo(new ChunkEntry(1, 0L, "foo").hashCode());
    }

    private ByteBuf[] createChunkedMessage(int messageSize, int maxChunkSize) {
        return createChunkedMessage(messageSize, maxChunkSize, generateMessageId());
    }

    private ByteBuf[] createChunkedMessage(int messageSize, int maxChunkSize, byte[] messageId) {
        // partially copied from GelfClient (can't use here, because it uses netty4)

        int sequenceCount = (messageSize / maxChunkSize);

        // Check if we have to add another chunk due to integer division.
        if ((messageSize % maxChunkSize) != 0) {
            sequenceCount++;
        }

        final ByteBuf[] buffers = new ByteBuf[sequenceCount];
        for (int sequenceNumber = 0; sequenceNumber < sequenceCount; sequenceNumber++) {
            // fake payload, we don't care about actually parsing it in this test
            int payloadSize = maxChunkSize;
            // correctly size the last chunk
            if (sequenceNumber + 1 == sequenceCount) {
                payloadSize = (messageSize % maxChunkSize);
            }

            buffers[sequenceNumber] = createChunk(messageId, (byte) sequenceNumber, (byte) sequenceCount,
                    new byte[payloadSize]);
        }

        return buffers;
    }

    private ByteBuf createChunk(byte[] messageId, byte sequenceNumber, byte sequenceCount, byte[] payload) {
        final ByteBuf channelBuffer = Unpooled.buffer(payload.length + 12);

        channelBuffer.writeBytes(CHUNK_MAGIC_BYTES);
        channelBuffer.writeBytes(messageId);
        channelBuffer.writeByte(sequenceNumber);
        channelBuffer.writeByte(sequenceCount);
        channelBuffer.writeBytes(payload);

        return channelBuffer;
    }

    private byte[] generateMessageId(int id) {
        final ByteBuf messageId = Unpooled.buffer(8);

        // 4 bytes of current time.
        messageId.writeInt((int) System.currentTimeMillis());
        messageId.writeInt(id);

        return messageId.array();
    }

    private byte[] generateMessageId() {
        return generateMessageId(0);
    }

    public static long counterValueNamed(MetricRegistry metricRegistry, String name) {
        return metricRegistry.getCounters(new SingleNameMatcher(name)).get(name).getCount();
    }

    private static class SingleNameMatcher implements MetricFilter {
        private final String metricName;

        public SingleNameMatcher(String metricName) {
            this.metricName = metricName;
        }

        @Override
        public boolean matches(String name, Metric metric) {
            return metricName.equals(name);
        }
    }
}