/* * Copyright 2011-2014 Proofpoint, Inc. * * 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.proofpoint.event.collector; import com.google.common.collect.ImmutableList; import com.proofpoint.event.collector.BatchProcessor.BatchHandler; import com.proofpoint.event.collector.queue.Queue; import com.proofpoint.event.collector.queue.QueueFactory; import com.proofpoint.reporting.ReportExporter; import com.proofpoint.testing.FileUtils; import com.proofpoint.units.Duration; import org.joda.time.DateTime; import org.mockito.Matchers; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; import java.io.File; import java.io.IOException; import java.util.Collections; import java.util.List; import java.util.UUID; import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; import static org.mockito.Matchers.anyInt; import static org.mockito.Matchers.anyString; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.timeout; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; import static org.testng.Assert.fail; public class TestAsyncBatchProcessor { private static final String DATA_DIRECTORY = "target/queue/events"; private BatchHandler<Event> handler; private BatchProcessorConfig config; private Queue<Event> queue; private QueueFactory<Event> queueFactory; private ReportExporter mockReporter; @BeforeMethod public void setup() throws IOException { FileUtils.deleteRecursively(new File(DATA_DIRECTORY)); config = new BatchProcessorConfig().setDataDirectory(DATA_DIRECTORY).setThrottleTime(new Duration(10, TimeUnit.MILLISECONDS)); //noinspection unchecked handler = mock(BatchHandler.class); mockReporter = mock(ReportExporter.class); queueFactory = new QueueFactory<>(config, mockReporter); queue = queueFactory.create("queue"); } @AfterMethod public void tearDown() throws IOException { queue.close(); } @Test(expectedExceptions = NullPointerException.class, expectedExceptionsMessageRegExp = "name is null") public void testConstructorNullName() throws IOException { new AsyncBatchProcessor<>(null, handler, new BatchProcessorConfig(), queueFactory); } @Test(expectedExceptions = NullPointerException.class, expectedExceptionsMessageRegExp = "handler is null") public void testConstructorNullHandler() throws IOException { new AsyncBatchProcessor<>("name", null, new BatchProcessorConfig(), queueFactory); } @Test(expectedExceptions = NullPointerException.class, expectedExceptionsMessageRegExp = "config is null") public void testConstructorNullConfig() throws IOException { new AsyncBatchProcessor<>("name", handler, null, queueFactory); } @Test(expectedExceptions = NullPointerException.class, expectedExceptionsMessageRegExp = "queueFactory is null") public void testConstructorNullQueueFactory() throws IOException { new AsyncBatchProcessor<>("name", handler, new BatchProcessorConfig(), null); } @Test(expectedExceptions = NullPointerException.class, expectedExceptionsMessageRegExp = "throttle time is null") public void testConstructorNullThrottleTime() throws IOException { new AsyncBatchProcessor<>("name", handler, new BatchProcessorConfig().setThrottleTime(null), queueFactory); } @Test public void testEnqueue() throws Exception { config.setMaxBatchSize(100).setQueueSize(100); BatchProcessor<Event> processor = new AsyncBatchProcessor<>( "foo", handler, config, queueFactory); processor.start(); try { processor.put(event("foo")); processor.put(event("foo")); processor.put(event("foo")); } finally { processor.stop(); } } @Test public void testFullQueue() throws Exception { config.setMaxBatchSize(100).setQueueSize(1); BlockingBatchHandler blockingHandler = blockingHandler(); BatchProcessor<Event> processor = new AsyncBatchProcessor<>( "foo", blockingHandler, config, new QueueFactory<Event>(config, mockReporter)); processor.start(); blockingHandler.lock(); try { // This will be processed, and its processing will block the handler processor.put(event("foo")); // Wait for the handler to pick up the item from the queue assertTrue(blockingHandler.waitForProcessor(10)); assertEquals(blockingHandler.getDroppedEntries(), 0); // This will remain in the queue because the processing // thread has not yet been resumed processor.put(event("foo")); assertEquals(blockingHandler.getDroppedEntries(), 0); // The queue is now full, this message will be dropped. processor.put(event("foo")); assertEquals(blockingHandler.getDroppedEntries(), 1); } finally { blockingHandler.resumeProcessor(); blockingHandler.unlock(); processor.stop(); queue.close(); } } @Test public void testContinueOnHandlerException() throws InterruptedException, IOException { config.setMaxBatchSize(100).setQueueSize(100); BlockingBatchHandler blockingHandler = blockingHandlerThatThrowsException(new RuntimeException("Expected Exception")); BatchProcessor<Event> processor = new AsyncBatchProcessor<>( "foo", blockingHandler, config, new QueueFactory<Event>(config, mockReporter)); processor.start(); blockingHandler.lock(); try { processor.put(event("foo")); assertTrue(blockingHandler.waitForProcessor(10)); } finally { blockingHandler.resumeProcessor(); blockingHandler.unlock(); } // When the processor unblocks, an exception (in its thread) will be thrown. // This should not affect the processor. blockingHandler.lock(); try { processor.put(event("bar")); assertTrue(blockingHandler.waitForProcessor(10)); } finally { blockingHandler.resumeProcessor(); blockingHandler.unlock(); processor.stop(); } } @Test public void testStopsWhenStopCalled() throws InterruptedException, IOException { config.setMaxBatchSize(100).setQueueSize(100); BlockingBatchHandler blockingHandler = blockingHandler(); BatchProcessor<Event> processor = new AsyncBatchProcessor<>( "foo", blockingHandler, config, new QueueFactory<Event>(config, mockReporter)); processor.start(); blockingHandler.lock(); try { processor.put(event("foo")); assertTrue(blockingHandler.waitForProcessor(10)); // The processor hasn't been resumed. Stop it! processor.stop(); } finally { blockingHandler.unlock(); } try { processor.put(event("bar")); fail(); } catch (IllegalStateException ex) { assertEquals(ex.getMessage(), "Processor is not running"); } } @Test public void testIgnoresExceptionsInHandler() throws InterruptedException, IOException { config.setMaxBatchSize(100).setQueueSize(100); BlockingBatchHandler blockingHandler = blockingHandlerThatThrowsException(new NullPointerException("Expected Exception")); BatchProcessor<Event> processor = new AsyncBatchProcessor<>( "foo", blockingHandler, config, new QueueFactory<Event>(config, mockReporter)); processor.start(); blockingHandler.lock(); try { processor.put(event("foo")); assertTrue(blockingHandler.waitForProcessor(10)); blockingHandler.resumeProcessor(); } finally { blockingHandler.unlock(); } blockingHandler.lock(); try { processor.put(event("foo")); assertTrue(blockingHandler.waitForProcessor(10)); blockingHandler.resumeProcessor(); } finally { blockingHandler.unlock(); processor.stop(); } } @Test @SuppressWarnings("unchecked") public void test_EmptyEntries_Ignored() throws Exception { config.setMaxBatchSize(100).setQueueSize(1).setThrottleTime(new Duration(100, TimeUnit.MILLISECONDS)); List<Event> events = Collections.emptyList(); BatchHandler<Event> mockHandler = mock(BatchHandler.class); Queue<Event> mockQueue = mock(Queue.class); when(mockQueue.dequeue(anyInt())).thenReturn(events); QueueFactory<Event> mockQueueFactory = mock(QueueFactory.class); when(mockQueueFactory.create(anyString())).thenReturn(mockQueue); BatchProcessor<Event> processor = new AsyncBatchProcessor<>("foo", mockHandler, config, mockQueueFactory); processor.start(); Thread.sleep(300); verify(mockHandler, never()).processBatch(Matchers.<List<Event>>any()); } @SuppressWarnings("unchecked") @Test public void test_failedEntries_GetEnqueued() throws Exception { config.setMaxBatchSize(100).setQueueSize(1).setThrottleTime(new Duration(100, TimeUnit.MILLISECONDS)); List<Event> events = ImmutableList.of(event("typeA1"), event("typeA2")); BatchHandler<Event> mockHandler = mock(BatchHandler.class); when(mockHandler.processBatch(events)).thenReturn(false); Queue<Event> mockQueue = mock(Queue.class); when(mockQueue.dequeue(anyInt())).thenReturn(events); when(mockQueue.enqueueAllOrDrop(events)).thenReturn(events); QueueFactory<Event> mockQueueFactory = mock(QueueFactory.class); when(mockQueueFactory.create(anyString())).thenReturn(mockQueue); BatchProcessor<Event> processor = new AsyncBatchProcessor<>("foo", mockHandler, config, mockQueueFactory); processor.start(); verify(mockQueue, timeout(200).times(1)).enqueueAllOrDrop(events); } private static Event event(String type) { return new Event(type, UUID.randomUUID().toString(), "localhost", DateTime.now(), Collections.<String, Object>emptyMap()); } private static BlockingBatchHandler blockingHandler() { return new BlockingBatchHandler(new Runnable() { @Override public void run() { } }); } private static BlockingBatchHandler blockingHandlerThatThrowsException(final RuntimeException exception) { return new BlockingBatchHandler(new Runnable() { @Override public void run() { throw exception; } }); } private static class BlockingBatchHandler implements BatchHandler<Event> { private final Lock lock = new ReentrantLock(); private final Condition externalCondition = lock.newCondition(); private final Condition internalCondition = lock.newCondition(); private final Runnable onProcess; private long droppedEntries = 0; public BlockingBatchHandler(Runnable onProcess) { this.onProcess = onProcess; } @Override public boolean processBatch(List<Event> entries) { // Wait for the right time to run lock.lock(); try { // Signal that we've started running if (entries.size() > 0) { externalCondition.signal(); try { // Block internalCondition.await(); } catch (InterruptedException ignored) { } } onProcess.run(); } finally { lock.unlock(); } return true; } @Override public void notifyEntriesDropped(int count) { droppedEntries += count; } public void lock() { lock.lock(); } public void unlock() { lock.unlock(); } public boolean waitForProcessor(long seconds) throws InterruptedException { return externalCondition.await(seconds, TimeUnit.SECONDS); } public void resumeProcessor() { internalCondition.signal(); } public long getDroppedEntries() { return droppedEntries; } } }