Updates to the S3 sink to speed up the unit test time. There are a fe…

…w major changes - use the Duration class instead of a nebulous long to have millisecond options and clarity; inject the retry sleep time so that the tests can sleep for shorter time; using mocking where possible to avoid unnecessary sleeps. (#3203)

Signed-off-by: David Venable <[email protected]>

data-prepper-plugins/s3-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkServiceIT.java

@@ -242,7 +242,7 @@ void verify_flushed_records_into_s3_bucketNewLine_with_compression() throws IOEx
 
     private S3SinkService createObjectUnderTest() {
         OutputCodecContext codecContext = new OutputCodecContext("Tag", Collections.emptyList(), Collections.emptyList());
-        return new S3SinkService(s3SinkConfig, bufferFactory, codec, codecContext, s3Client, keyGenerator, pluginMetrics);
+        return new S3SinkService(s3SinkConfig, bufferFactory, codec, codecContext, s3Client, keyGenerator, Duration.ofSeconds(5), pluginMetrics);
     }
 
     private int gets3ObjectCount() {

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/S3Sink.java

@@ -27,6 +27,7 @@
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.services.s3.S3Client;
 
+import java.time.Duration;
 import java.util.Collection;
 
 /**
@@ -75,7 +76,7 @@ public S3Sink(final PluginSetting pluginSetting,
 
         S3OutputCodecContext s3OutputCodecContext = new S3OutputCodecContext(OutputCodecContext.fromSinkContext(sinkContext), compressionOption);
 
-        s3SinkService = new S3SinkService(s3SinkConfig, bufferFactory, codec, s3OutputCodecContext, s3Client, keyGenerator, pluginMetrics);
+        s3SinkService = new S3SinkService(s3SinkConfig, bufferFactory, codec, s3OutputCodecContext, s3Client, keyGenerator, Duration.ofSeconds(5), pluginMetrics);
     }
 
     @Override

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkService.java

@@ -23,6 +23,7 @@
 import software.amazon.awssdk.services.s3.S3Client;
 
 import java.io.IOException;
+import java.time.Duration;
 import java.util.Collection;
 import java.util.LinkedList;
 import java.util.concurrent.locks.Lock;
@@ -49,7 +50,7 @@ public class S3SinkService {
     private Buffer currentBuffer;
     private final int maxEvents;
     private final ByteCount maxBytes;
-    private final long maxCollectionDuration;
+    private final Duration maxCollectionDuration;
     private final String bucket;
     private final int maxRetries;
     private final Counter objectsSucceededCounter;
@@ -59,6 +60,7 @@ public class S3SinkService {
     private final DistributionSummary s3ObjectSizeSummary;
     private final OutputCodecContext codecContext;
     private final KeyGenerator keyGenerator;
+    private final Duration retrySleepTime;
 
     /**
      * @param s3SinkConfig  s3 sink related configuration.
@@ -68,20 +70,22 @@ public class S3SinkService {
      * @param pluginMetrics metrics.
      */
     public S3SinkService(final S3SinkConfig s3SinkConfig, final BufferFactory bufferFactory,
-                         final OutputCodec codec, final OutputCodecContext codecContext, final S3Client s3Client, final KeyGenerator keyGenerator, final PluginMetrics pluginMetrics) {
+                         final OutputCodec codec, final OutputCodecContext codecContext, final S3Client s3Client, final KeyGenerator keyGenerator,
+                         final Duration retrySleepTime, final PluginMetrics pluginMetrics) {
         this.s3SinkConfig = s3SinkConfig;
         this.bufferFactory = bufferFactory;
         this.codec = codec;
         this.s3Client = s3Client;
         this.codecContext = codecContext;
         this.keyGenerator = keyGenerator;
+        this.retrySleepTime = retrySleepTime;
         reentrantLock = new ReentrantLock();
 
         bufferedEventHandles = new LinkedList<>();
 
         maxEvents = s3SinkConfig.getThresholdOptions().getEventCount();
         maxBytes = s3SinkConfig.getThresholdOptions().getMaximumSize();
-        maxCollectionDuration = s3SinkConfig.getThresholdOptions().getEventCollectTimeOut().getSeconds();
+        maxCollectionDuration = s3SinkConfig.getThresholdOptions().getEventCollectTimeOut();
 
         bucket = s3SinkConfig.getBucketName();
         maxRetries = s3SinkConfig.getMaxUploadRetries();
@@ -191,7 +195,7 @@ protected boolean retryFlushToS3(final Buffer currentBuffer, final String s3Key)
                 }
 
                 try {
-                    Thread.sleep(5000);
+                    Thread.sleep(retrySleepTime.toMillis());
                 } catch (final InterruptedException ex) {
                     LOG.warn("Interrupted while backing off before retrying S3 upload", ex);
                 }

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/ThresholdCheck.java

@@ -8,6 +8,8 @@
 import org.opensearch.dataprepper.model.types.ByteCount;
 import org.opensearch.dataprepper.plugins.sink.s3.accumulator.Buffer;
 
+import java.time.Duration;
+
 /**
  * Check threshold limits.
  */
@@ -24,13 +26,13 @@ private ThresholdCheck() {
      * @param maxCollectionDuration maximum event collection duration provided by user as threshold.
      * @return boolean value whether the threshold are met.
      */
-    public static boolean checkThresholdExceed(final Buffer currentBuffer, final int maxEvents, final ByteCount maxBytes, final long maxCollectionDuration) {
+    public static boolean checkThresholdExceed(final Buffer currentBuffer, final int maxEvents, final ByteCount maxBytes, final Duration maxCollectionDuration) {
         if (maxEvents > 0) {
             return currentBuffer.getEventCount() + 1 > maxEvents ||
-                    currentBuffer.getDuration() > maxCollectionDuration ||
+                    currentBuffer.getDuration().compareTo(maxCollectionDuration) > 0 ||
                     currentBuffer.getSize() > maxBytes.getBytes();
         } else {
-            return currentBuffer.getDuration() > maxCollectionDuration ||
+            return currentBuffer.getDuration().compareTo(maxCollectionDuration) > 0 ||
                     currentBuffer.getSize() > maxBytes.getBytes();
         }
     }

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/accumulator/Buffer.java

@@ -6,6 +6,7 @@
 package org.opensearch.dataprepper.plugins.sink.s3.accumulator;
 
 import java.io.OutputStream;
+import java.time.Duration;
 
 /**
  * A buffer can hold data before flushing it to S3.
@@ -18,7 +19,7 @@ public interface Buffer {
     long getSize();
     int getEventCount();
 
-    long getDuration();
+    Duration getDuration();
 
     void flushToS3();
 

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/accumulator/CompressionBuffer.java

@@ -9,6 +9,7 @@
 
 import java.io.IOException;
 import java.io.OutputStream;
+import java.time.Duration;
 import java.util.Objects;
 
 class CompressionBuffer implements Buffer {
@@ -32,7 +33,7 @@ public int getEventCount() {
     }
 
     @Override
-    public long getDuration() {
+    public Duration getDuration() {
         return innerBuffer.getDuration();
     }
 

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/accumulator/InMemoryBuffer.java

@@ -11,6 +11,7 @@
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import java.io.ByteArrayOutputStream;
 import java.io.OutputStream;
+import java.time.Duration;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 
@@ -50,8 +51,8 @@ public int getEventCount() {
         return eventCount;
     }
 
-    public long getDuration() {
-        return watch.getTime(TimeUnit.SECONDS);
+    public Duration getDuration() {
+        return Duration.ofMillis(watch.getTime(TimeUnit.MILLISECONDS));
     }
 
     /**

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/accumulator/LocalFileBuffer.java

@@ -19,6 +19,7 @@
 import java.io.OutputStream;
 import java.nio.file.Files;
 import java.nio.file.Paths;
+import java.time.Duration;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 
@@ -68,8 +69,8 @@ public int getEventCount() {
     }
 
     @Override
-    public long getDuration(){
-        return watch.getTime(TimeUnit.SECONDS);
+    public Duration getDuration(){
+        return Duration.ofMillis(watch.getTime(TimeUnit.MILLISECONDS));
     }
 
     /**

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/accumulator/MultipartBuffer.java

@@ -9,6 +9,7 @@
 import org.opensearch.dataprepper.plugins.codec.parquet.S3OutputStream;
 
 import java.io.IOException;
+import java.time.Duration;
 import java.util.concurrent.TimeUnit;
 
 public class MultipartBuffer implements Buffer {
@@ -40,8 +41,8 @@ public int getEventCount() {
         return eventCount;
     }
 
-    public long getDuration() {
-        return watch.getTime(TimeUnit.SECONDS);
+    public Duration getDuration() {
+        return Duration.ofMillis(watch.getTime(TimeUnit.MILLISECONDS));
     }
 
     /**

data-prepper-plugins/s3-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkServiceTest.java

@@ -133,7 +133,7 @@ void setUp() {
     }
 
     private S3SinkService createObjectUnderTest() {
-        return new S3SinkService(s3SinkConfig, bufferFactory, codec, codecContext, s3Client, keyGenerator, pluginMetrics);
+        return new S3SinkService(s3SinkConfig, bufferFactory, codec, codecContext, s3Client, keyGenerator, Duration.ofMillis(100), pluginMetrics);
     }
 
     @Test