/* * 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.base.Charsets; import com.google.common.base.MoreObjects; import com.google.common.collect.Maps; import com.google.common.collect.Sets; import com.google.common.io.CharStreams; import com.google.common.io.Files; import com.google.common.io.Resources; import com.google.common.util.concurrent.ThreadFactoryBuilder; import com.proofpoint.event.collector.combiner.CombinedStoredObject; import com.proofpoint.event.collector.combiner.StorageSystem; import com.proofpoint.event.collector.combiner.StoredObject; import com.proofpoint.reporting.testing.TestingReportCollectionFactory; import com.proofpoint.stats.SparseCounterStat; import com.proofpoint.stats.SparseTimeStat; import com.proofpoint.testing.SerialScheduledExecutorService; import com.proofpoint.units.Duration; import org.iq80.snappy.SnappyInputStream; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; import java.io.File; import java.io.FileInputStream; import java.io.IOException; import java.io.InputStreamReader; import java.net.URI; import java.net.URISyntaxException; import java.util.List; import java.util.Map; import java.util.Set; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import static com.proofpoint.event.collector.S3UploaderStats.FileProcessedStatus.CORRUPT; import static com.proofpoint.event.collector.S3UploaderStats.FileProcessedStatus.UPLOADED; import static com.proofpoint.event.collector.S3UploaderStats.FileUploadStatus.FAILURE; import static com.proofpoint.event.collector.S3UploaderStats.FileUploadStatus.SUCCESS; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; public class TestS3Uploader { private static final String ARBITRARY_EVENT_TYPE = "TapPrsMessage"; private static final String UNKNOWN_EVENT_TYPE = "unknown"; private File tempStageDir; private S3Uploader uploader; private DummyStorageSystem storageSystem; private ServerConfig serverConfig; private SerialScheduledExecutorService executor; private TestingReportCollectionFactory testingReportCollectionFactory; private S3UploaderStats s3UploaderStats; @BeforeMethod public void setup() { storageSystem = new DummyStorageSystem(); tempStageDir = Files.createTempDir(); tempStageDir.deleteOnExit(); serverConfig = new ServerConfig() .setLocalStagingDirectory(tempStageDir) .setS3StagingLocation("s3://fake-location") .setRetryPeriod(new Duration(3, TimeUnit.MINUTES)) .setRetryDelay(new Duration(1, TimeUnit.MINUTES)); executor = new SerialScheduledExecutorService(); testingReportCollectionFactory = new TestingReportCollectionFactory(); s3UploaderStats = testingReportCollectionFactory.createReportCollection(S3UploaderStats.class); } @Test(expectedExceptions = NullPointerException.class, expectedExceptionsMessageRegExp = "s3UploaderStats is null") public void testConstructorNullS3UploaderStats() { new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), executor, executor, null); } @Test public void testSuccessOnFailedDirExists() { new File(tempStageDir.getPath(), "failed").mkdir(); uploader = new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), executor, executor, s3UploaderStats); } @Test(expectedExceptions = IllegalArgumentException.class) public void testFailureOnFailedFileExists() throws IOException { new File(tempStageDir.getPath(), "failed").createNewFile(); uploader = new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), executor, executor, s3UploaderStats); } @Test public void testSuccessOnRetryDirExists() { new File(tempStageDir.getPath(), "retry").mkdir(); uploader = new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), executor, executor, s3UploaderStats); } @Test(expectedExceptions = IllegalArgumentException.class) public void testFailureOnRetryFileExists() throws IOException { new File(tempStageDir.getPath(), "retry").createNewFile(); uploader = new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), executor, executor, s3UploaderStats); } @Test public void testSuccessOnDirectoriesInStaging() throws IOException { new File(tempStageDir.getPath(), "directory").mkdir(); uploader = new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), executor, executor, s3UploaderStats); uploader.start(); } @Test public void testUploadPendingFiles() throws Exception { File pendingFile = new File(tempStageDir, "pending.json.snappy"); Files.copy(new File(Resources.getResource("pending.json.snappy").toURI()), pendingFile); assertTrue(pendingFile.exists()); uploader = new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), executor, executor, s3UploaderStats); uploader.start(); assertFalse(pendingFile.exists()); assertTrue(storageSystem.hasReceivedFile(pendingFile)); S3UploaderStats s3UploaderStatsArgumentVerifier = testingReportCollectionFactory.getArgumentVerifier(S3UploaderStats.class); verify(s3UploaderStatsArgumentVerifier).processedFiles(ARBITRARY_EVENT_TYPE, UPLOADED); verify(s3UploaderStatsArgumentVerifier).uploadAttempts(ARBITRARY_EVENT_TYPE, SUCCESS); verify(s3UploaderStatsArgumentVerifier).processedTime(ARBITRARY_EVENT_TYPE); verifyNoMoreInteractions(s3UploaderStatsArgumentVerifier); S3UploaderStats s3UploaderStatsReportCollection = testingReportCollectionFactory.getReportCollection(S3UploaderStats.class); SparseCounterStat processedFilesCounterStat = s3UploaderStatsReportCollection.processedFiles(ARBITRARY_EVENT_TYPE, UPLOADED); SparseCounterStat uploadAttemptsCounterStat = s3UploaderStatsReportCollection.uploadAttempts(ARBITRARY_EVENT_TYPE, SUCCESS); SparseTimeStat processedTimeTimerStat = s3UploaderStatsReportCollection.processedTime(ARBITRARY_EVENT_TYPE); verify(processedFilesCounterStat).add(1); verify(uploadAttemptsCounterStat).add(1); verify(processedTimeTimerStat).time(); verifyNoMoreInteractions(processedFilesCounterStat); verifyNoMoreInteractions(uploadAttemptsCounterStat); verifyNoMoreInteractions(processedTimeTimerStat); } @Test public void testUploadPendingFilesFailure() throws Exception { File invalidJsonFile = new File(tempStageDir, "invalidjson.snappy"); Files.copy(new File(Resources.getResource("invalidjson.snappy").toURI()), invalidJsonFile); assertTrue(invalidJsonFile.exists()); uploader = new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), executor, executor, s3UploaderStats); uploader.start(); assertFalse(invalidJsonFile.exists()); assertTrue(new File(tempStageDir.getPath() + "/failed", invalidJsonFile.getName()).exists()); assertFalse(storageSystem.hasReceivedFile(invalidJsonFile)); S3UploaderStats s3UploaderStatsArgumentVerifier = testingReportCollectionFactory.getArgumentVerifier(S3UploaderStats.class); verify(s3UploaderStatsArgumentVerifier).processedFiles(UNKNOWN_EVENT_TYPE, CORRUPT); verifyNoMoreInteractions(s3UploaderStatsArgumentVerifier); SparseCounterStat processedFilesCounterStat = testingReportCollectionFactory.getReportCollection(S3UploaderStats.class).processedFiles(UNKNOWN_EVENT_TYPE, CORRUPT); verify(processedFilesCounterStat).add(1); verifyNoMoreInteractions(processedFilesCounterStat); } @Test public void testFailsOnVerifyFile() throws Exception { File invalidJsonFile = new File(tempStageDir, "invalidjson2.snappy"); Files.copy(new File(Resources.getResource("invalidjson2.snappy").toURI()), invalidJsonFile); assertTrue(invalidJsonFile.exists()); uploader = new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), executor, executor, s3UploaderStats); uploader.start(); assertFalse(invalidJsonFile.exists()); assertTrue(new File(tempStageDir.getPath() + "/failed", invalidJsonFile.getName()).exists()); assertFalse(storageSystem.hasReceivedFile(invalidJsonFile)); S3UploaderStats s3UploaderStatsArgumentVerifier = testingReportCollectionFactory.getArgumentVerifier(S3UploaderStats.class); verify(s3UploaderStatsArgumentVerifier).processedFiles(UNKNOWN_EVENT_TYPE, CORRUPT); verifyNoMoreInteractions(s3UploaderStatsArgumentVerifier); SparseCounterStat processedFilesCounterStat = testingReportCollectionFactory.getReportCollection(S3UploaderStats.class).processedFiles(UNKNOWN_EVENT_TYPE, CORRUPT); verify(processedFilesCounterStat).add(1); verifyNoMoreInteractions(processedFilesCounterStat); } @Test public void testRetryUntilSuccess() throws Exception { storageSystem.succeedOnAttempt(3); File pendingFile = new File(tempStageDir, "pending.json.snappy"); Files.copy(new File(Resources.getResource("pending.json.snappy").toURI()), pendingFile); String retryDir = tempStageDir.getPath() + "/retry"; assertEquals(storageSystem.getAttempts(pendingFile), 0); //attempt 1 to upload from staging directory fails and file is moved to retry directory assertTrue(pendingFile.exists()); uploader = new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), executor, executor, s3UploaderStats); uploader.start(); assertFalse(pendingFile.exists()); assertTrue(new File(retryDir, pendingFile.getName()).exists()); assertFalse(storageSystem.hasReceivedFile(pendingFile)); assertEquals(storageSystem.getAttempts(pendingFile), 1); //attempt 2: file is moved to staging from retry directory, and fails and hence gets back to retry directory executor.elapseTime(1, TimeUnit.MINUTES); assertFalse(pendingFile.exists()); assertTrue(new File(retryDir, pendingFile.getName()).exists()); assertFalse(storageSystem.hasReceivedFile(pendingFile)); assertEquals(storageSystem.getAttempts(pendingFile), 2); //retryExecutor hasn't run again executor.elapseTime(1, TimeUnit.MINUTES); assertFalse(pendingFile.exists()); assertTrue(new File(retryDir, pendingFile.getName()).exists()); assertFalse(storageSystem.hasReceivedFile(pendingFile)); assertEquals(storageSystem.getAttempts(pendingFile), 2); //attempt 3: file is moved to staging from retry directory, succeeds and hence is deleted from local directories executor.elapseTime(2, TimeUnit.MINUTES); assertFalse(pendingFile.exists()); assertFalse(new File(retryDir, pendingFile.getName()).exists()); assertTrue(storageSystem.hasReceivedFile(pendingFile)); assertEquals(storageSystem.getAttempts(pendingFile), 3); } @Test public void testPulseMetricsForUploadAttempts() throws Exception { storageSystem.succeedOnAttempt(2); File pendingFile = new File(tempStageDir, "pending.json.snappy"); Files.copy(new File(Resources.getResource("pending.json.snappy").toURI()), pendingFile); //attempt 1 to upload from staging directory fails and file is moved to retry directory uploader = new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), executor, executor, s3UploaderStats); uploader.start(); S3UploaderStats s3UploaderStatsArgumentVerifier = testingReportCollectionFactory.getArgumentVerifier(S3UploaderStats.class); verify(s3UploaderStatsArgumentVerifier).uploadAttempts(ARBITRARY_EVENT_TYPE, FAILURE); verify(s3UploaderStatsArgumentVerifier).processedTime(ARBITRARY_EVENT_TYPE); verifyNoMoreInteractions(s3UploaderStatsArgumentVerifier); S3UploaderStats s3UploaderStatsReportCollection = testingReportCollectionFactory.getReportCollection(S3UploaderStats.class); SparseCounterStat uploadAttemptsCounterStat = s3UploaderStatsReportCollection.uploadAttempts(ARBITRARY_EVENT_TYPE, FAILURE); SparseTimeStat processedTimeTimerStat = s3UploaderStatsReportCollection.processedTime(ARBITRARY_EVENT_TYPE); verify(uploadAttemptsCounterStat).add(1); verify(processedTimeTimerStat).time(); verifyNoMoreInteractions(uploadAttemptsCounterStat); verifyNoMoreInteractions(processedTimeTimerStat); } @Test public void testInvalidFilesInRetryDirectory() throws Exception { uploader = new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), executor, executor, s3UploaderStats); uploader.start(); executor.elapseTime(1, TimeUnit.MINUTES); String retryDir = tempStageDir.getPath() + "/retry"; String failedDir = tempStageDir.getPath() + "/failed"; File invalidJsonFile = new File(retryDir, "invalidjson.snappy"); Files.copy(new File(Resources.getResource("invalidjson.snappy").toURI()), invalidJsonFile); File invalidFile = new File(retryDir, "invalidFile.snappy"); Files.copy(new File(Resources.getResource("invalidjson2.snappy").toURI()), invalidFile); File directory = new File(retryDir, "subdir"); directory.mkdir(); assertTrue(invalidJsonFile.exists()); assertTrue(invalidFile.exists()); assertTrue(directory.exists()); executor.elapseTime(3, TimeUnit.MINUTES); assertTrue(new File(failedDir, invalidJsonFile.getName()).exists()); assertFalse(new File(retryDir, invalidJsonFile.getName()).exists()); assertFalse(new File(tempStageDir.getPath(), invalidJsonFile.getName()).exists()); assertTrue(new File(failedDir, invalidFile.getName()).exists()); assertFalse(new File(retryDir, invalidFile.getName()).exists()); assertFalse(new File(tempStageDir.getPath(), invalidFile.getName()).exists()); assertTrue(directory.exists()); } @Test public void test_invalidSnappyFileGetsDecompressed() throws Exception { String json = unsnappy("invalidjson.snappy"); assertNotNull(json); } private String unsnappy(String fileName) throws URISyntaxException, IOException { File file = new File(Resources.getResource(fileName).toURI()); SnappyInputStream input = new SnappyInputStream(new FileInputStream(file)); return CharStreams.toString(new InputStreamReader(input, Charsets.UTF_8)); } @Test public void test_destroy_waitsForTasksToComplete() throws Exception { ExecutorService uploadExecutor = Executors.newFixedThreadPool(5, new ThreadFactoryBuilder().setNameFormat("S3Uploader-%s").build()); ScheduledExecutorService retryExecutor = Executors.newSingleThreadScheduledExecutor(); uploader = new S3Uploader(storageSystem, serverConfig, new EventPartitioner(), uploadExecutor, retryExecutor, s3UploaderStats); uploader.start(); final AtomicInteger counter = new AtomicInteger(0); executor.submit(new LoopingIncrementor(counter)); retryExecutor.submit(new LoopingIncrementor(counter)); uploader.destroy(); assertEquals(counter.get(), 2); } private class LoopingIncrementor implements Runnable { private final AtomicInteger counter; private LoopingIncrementor(AtomicInteger counter) { this.counter = counter; } @Override public void run() { try { for (int i = 0; i < 10; i++) { Thread.yield(); Thread.sleep(50); } counter.incrementAndGet(); } catch (InterruptedException e) { throw new RuntimeException(e); } } } private class DummyStorageSystem implements StorageSystem { private Set<String> receivedFiles = Sets.newHashSet(); private Map<String, Integer> retryCounts = Maps.newHashMap(); private int succeedOnAttempt = 1; public boolean hasReceivedFile(File file) { return receivedFiles.contains(file.getName()); } public void succeedOnAttempt(int attempt) { succeedOnAttempt = attempt; } public int getAttempts(File file) { return MoreObjects.firstNonNull(retryCounts.get(file.getName()), 0); } @Override public StoredObject putObject(URI location, File source) { int pastAttempts = getAttempts(source); int currentAttempt = pastAttempts + 1; retryCounts.put(source.getName(), currentAttempt); if (currentAttempt == succeedOnAttempt) { receivedFiles.add(source.getName()); return new StoredObject(URI.create("s3://dummyUri/bucket/day/hour")); } throw new RuntimeException("Exception in DummyStorageSystem"); } @Override public List<URI> listDirectories(URI storageArea) { throw new UnsupportedOperationException("Not supported yet."); } @Override public List<StoredObject> listObjects(URI storageArea) { throw new UnsupportedOperationException("Not supported yet."); } @Override public StoredObject createCombinedObject(CombinedStoredObject combinedObject) { throw new UnsupportedOperationException("Not supported yet."); } } }